from contextlib import contextmanager
import numpy as np
import tensorflow as tf
from time import time
from analysis.util import atq
from io import BytesIO
import h5py
from tensorflow.python.keras.engine.saving import (
save_weights_to_hdf5_group,
load_weights_from_hdf5_group,
)
from evil import pin, ccall
from tensorflow.python.util import deprecation
deprecation._PRINT_DEPRECATION_WARNINGS = False
def last_map(func, obj, order=1):
return tf.concat(
[
obj[..., :-1],
tf.expand_dims(
last_map(func=func, obj=obj[..., -1], order=order - 1)
if 1 < order
else func(obj[..., -1]),
axis=-1,
),
],
axis=-1,
)
def tf_meanstd(val, **kwargs):
mean = tf.reduce_mean(val, **kwargs)
std = tf.reduce_mean(tf.square(val), **kwargs) - tf.square(mean)
return mean, std
# this is mostly unused ...
class FD(dict):
def add(self, data, dtype=None, shape=None, **kwargs):
if shape is True:
shape = data.shape
elif shape is None:
shape = (None,) + data.shape[1:]
if dtype is None:
dtype = data.dtype
ph = tf.placeholder(dtype, shape=shape, **kwargs)
self[ph] = data
return ph
class Chain(object):
def __init__(
self,
__name__,
loss,
sumnbs=(),
step=None,
opt=None,
train=None,
ema=(1, 2, 3, 4, 5),
**kwargs
):
if step is None:
step = tf.Variable(0, trainable=False, name="%s_step" % __name__)
if opt is None:
opt = tf.train.AdamOptimizer()
if train is None:
train = opt.minimize(loss, global_step=step)
reset_opt = tf.variables_initializer(opt.variables())
summaries = []
_ema_val = {}
_ema_out = SKDict()
kwargs["loss"] = loss
with tf.name_scope(__name__):
for key, val in kwargs.items():
if not isinstance(val, tf.Tensor):
continue
nd = val.shape.ndims
if nd > 1 and sumnbs is not True and val not in sumnbs:
continue
if val.dtype == tf.bool:
val = tf.reduce_mean(tf.cast(val, tf.float32))
if nd:
# summaries.append(tf.summary.histogram(key, val))
mean, std = tf_meanstd(val)
summaries.append(tf.summary.scalar(key + "_mean", mean))
summaries.append(tf.summary.scalar(key + "_std", std))
_ema_val[key + "_mean"] = mean
_ema_val[key + "_std"] = std
else:
summaries.append(tf.summary.scalar(key, val))
_ema_val[key] = val
if ema:
_ema_ops = []
reset_ema_to = tf.constant(0, tf.int32)
ema_step = tf.Variable(reset_ema_to, trainable=False, name="%s_ema_step" % __name__)
reset_ema = ema_step.initializer
for i in ema:
with tf.name_scope("%s_%d" % (__name__, i)):
ema = tf.train.ExponentialMovingAverage(
decay=1.0
- tf.maximum(
0.1 ** i,
tf.exp((np.log(0.1 ** i) / 10.0) * tf.cast(ema_step, tf.float32)),
)
)
_ema_ops.append(ema.apply(_ema_val.values()))
for key, val in _ema_val.items():
val = ema.average(val)
_ema_out[key, i] = val
summaries.append(tf.summary.scalar("%s" % key, val))
with tf.control_dependencies([train]):
with tf.control_dependencies(_ema_ops):
train = ema_step.assign_add(1)
ema = _ema_out
summaries = tf.summary.merge(summaries)
self.__dict__.update(kwargs)
pin(locals(), True, kwargs)
Chain.cc = ccall(Chain, True)
class Runner(object):
def __init__(self, sess, writer):
self.step = 0
pin(locals())
def __call__(self, *args, **kwargs):
return self.sess.run(*args, **kwargs)
def train(self, chain, extra=None, step=1, sumiv=10, sumskip=0, **kwargs):
if extra is None:
extra = chain.loss
for i in count():
for j in xrange(step):
self(chain.train, **kwargs)
sum = sumiv and not (i < sumskip or i % sumiv)
get = (chain.train, extra)
if sum:
get += (chain.summaries, chain.step)
out = self(get, **kwargs)
if sum:
self.writer.add_summary(*out[2:])
yield out[1]
def conf(self, *args, **kwargs):
return Confirmer(self.train(*args, **kwargs).next)
@classmethod
@contextmanager
def make(cls, path, gpuOpts={}, flush_secs=20):
with cls.make_session(gpuOpts=gpuOpts) as sess:
with tf.summary.FileWriter(
path, session=sess, flush_secs=flush_secs, graph=sess.graph
) as writer:
yield cls(sess, writer)
@classmethod
def make_session(cls, gpuOpts={}):
return tf.Session(config=tf.ConfigProto(gpu_options=tf.GPUOptions(**gpuOpts)))
class Confirmer(object):
last = np.inf
task = None
tdelay = 60
tbestbump = 10
def __init__(self, func):
self.reset.func = func
def __call__(self, target, its=np.inf, pre=0, **kwargs):
if not kwargs.setdefault("disable", not kwargs):
kwargs.setdefault("auto", self._prog)
assert 0 < target
tnext = 0
with atq(xrange(its), **kwargs) as tq:
for i in tq:
if self.task is not None:
if self.task.check_stop():
break
if tnext < time():
tnext = time() + self.tdelay
pTotal = 100.0 * i / its
self.task.set_progress_percentage(pTotal)
self.task.set_status_message(
"total=%.1f%% conf=%.1f%% best=%.2e"
% (pTotal, 100.0 * self.conf / target, self.best)
)
if not (self.conf < target):
break
self.step += 1
self.last = self.func()
if self.step < pre:
continue
elif self.best < self.last:
self.conf += 1
else:
self.on_best()
self.conf = 0
self.best = self.last
tnext -= self.tbestbump
return not (self.conf < target)
def on_best(self):
pass
@property
def reset(self):
self.last = self.best = np.inf
self.step = self.conf = 0
return self
@property
def rdlb(self):
# Relative Delta Last to Best
return np.float_(self.last - self.best) / self.best
def _prog(self, i, prefix=""):
return {
prefix + k: v
for k, v in dict(
step=self.step if self.step != i else None,
last=self.last if self.rdlb else None,
rdlb=self.rdlb or None,
conf=self.conf or None,
best=self.best,
).items()
if v is not None
}
def __repr__(self):
return "Confirmer(step=%d, best=%.3e, conf=%d)" % (self.step, self.best, self.conf)
class Bestie(object):
def __init__(self, model, path):
pin(locals())
def save(self):
self.model.save_weights(self.path)
def load(self):
self.model.load_weights(self.path)
def finish(self):
self.load()
self.model.save(self.path)
class BestieMemory(Bestie):
def __init__(self, *args, **kwargs):
super(BestieMemory, self).__init__(*args, **kwargs)
def save(self):
self.h5 = h5py.File(BytesIO())
save_weights_to_hdf5_group(self.h5, self.model.layers)
def load(self):
load_weights_from_hdf5_group(self.h5, self.model.layers)