def call(self, x):
output_mu = tf.matmul(x, self.kernel_1) + self.bias_1
output_sig = tf.matmul(x, self.kernel_2) + self.bias_2
output_sig = softplus(output_sig)
if self.count_data:
output_mu = softplus(output_mu)
return [output_mu, output_sig]
def _softplus(
self,
mu: tf.Tensor,
scale: tf.Tensor,
) -> typing.Tuple[tf.Tensor, tf.Tensor]:
"""
private util function that applies softplus transformation to various parameters
depending on type of data
"""
scale = softplus(scale)
if self._ts_obj.count_data:
mu = softplus(mu)
return mu, scale
def call(self, x):
X = [0] * 128
for i in range(self.n_kernels):
X[i] = softplus(
K.conv1d(x[:, :, i:i + 1], self.kernels[i], padding='same'))
K.bias_add(X[i], self.bias[i], data_format=self.data_format)
X = K.concatenate(X, axis=2)
return X
def call(self, x):
rank = len(x.shape)
X = [0] * x.shape[1]
for i in range(self.n_kernels):
# Broadcasting is required for the inputs.
X[i] = softplus(K.dot(x[:, i:i + 1, :], self.kernels[i]))
K.bias_add(X[i], self.bias[i], data_format='channels_last')
X = K.concatenate(X, axis=1)
return X
def __init__(self,
memory_size=128,
word_size=20,
num_reads=1,
num_writes=1,
memory_init_value=1e-6,
name='memory_access',
dtype=tf.float32):
"""Creates a MemoryAccess module.
Args:
memory_size: The number of memory slots (N in the DNC paper).
word_size: The width of each memory slot (W in the DNC paper)
num_reads: The number of read heads (R in the DNC paper).
num_writes: The number of write heads (fixed at 1 in the paper).
memory_init_value: The initial value for memory
name: The name of the module.
"""
super().__init__(name=name)
self._memory_size = memory_size
self._word_size = word_size
self._num_reads = num_reads
self._num_writes = num_writes
self._memory_init_value = tf.convert_to_tensor(memory_init_value, dtype)
self._write_content_weights_mod = addressing.CosineWeights(
num_writes, word_size, name='write_content_weights')
self._read_content_weights_mod = addressing.CosineWeights(
num_reads, word_size, name='read_content_weights')
self._linkage = addressing.TemporalLinkage(memory_size, num_writes)
self._freeness = addressing.Freeness(memory_size)
Dense = layers.Dense
sigmoid = activations.sigmoid
oneplus = lambda x: activations.softplus(x) + 1
self._layers = dict(
read_keys=Dense(self._num_reads*self._word_size, name='read_keys', dtype=dtype),
read_strengths=Dense(self._num_reads, name='read_strengths', activation=oneplus, dtype=dtype),
write_keys=Dense(self._num_writes*self._word_size, name='write_keys', dtype=dtype),
write_strengths=Dense(self._num_writes, name='write_strengths', activation=oneplus, dtype=dtype),
write_vectors=Dense(self._num_writes*self._word_size, name='write_vectors', dtype=dtype),
erase_vectors=Dense(self._num_writes*self._word_size, activation=sigmoid, name='erase_vectors', dtype=dtype),
free_gate=Dense(self._num_reads, activation=sigmoid, name='free_gate', dtype=dtype),
allocation_gate=Dense(self._num_writes, activation=sigmoid, name='allocation_gate', dtype=dtype),
write_gate=Dense(self._num_writes, activation=sigmoid, name='write_gate', dtype=dtype),
read_mode=Dense(self._num_reads * (2 * self._num_writes + 1),name='read_mode', dtype=dtype),
)
def __init__(self):
self.vae = AutoEncoder()
# def disc_loss_fn(true, out):
# tf.print(true)
# tf.print(out)
# return binary_crossentropy(true, out)
self.discriminator = self.create_discriminator()
self.discriminator.compile(Adam(lr=1e-4), loss='binary_crossentropy')
inp = Input(shape=(*imgdims, 1))
fake, repr = self.vae.model(inp)
self.discriminator.trainable = False
disc = self.discriminator(fake)
self.complete_model = Model(inp, [disc, repr], name='vae_gan')
self.complete_model.compile(
'adam',
loss=lambda true, out: binary_crossentropy(true[0], out[
0]) + softplus(img_variety(tf.argmax(out[1])) - 1000) / 100)
def _build_model(self, transfer_model=None, kernel_regularizer=None, kernel_initializer=None):
input_shape = (*tuple(self.config.glob.image_size), self.config.glob.image_n_chanel)
if self.config.model.transfer_model.exist:
transfer_model = transfer_models[self.config.model.transfer_model.name](
weights=self.config.model.transfer_model.weights,
include_top=False,
input_shape=input_shape
)
self.model = Sequential([
transfer_model,
GlobalAveragePooling2D()
])
n_layers = len(self.config.model.dense.units)
dense_units = self.config.model.dense.units
dense_activation = self.config.model.dense.activation
dropout_rates = self.config.model.dropout.rates
n_classes = self.config.glob.n_classes
if dense_activation == 'tanh_softplus':
dense_activation = lambda x: x * backend.tanh(activations.softplus(x))
if self.config.model.kernel_regularizer.exist:
kernel_regularizer_name = self.config.model.kernel_regularizer.name
alpha = self.config.model.kernel_regularizer.alpha
kernel_regularizer = regularizers[kernel_regularizer_name](alpha)
if self.config.model.kernel_initializer.exist:
kernel_initializer = self.config.model.kernel_initializer.name
for units, dropout_rate, _ in zip(dense_units, dropout_rates, range(n_layers)):
self.model.add(Dense(
units,
activation=dense_activation,
kernel_regularizer=kernel_regularizer,
kernel_initializer=kernel_initializer
))
self.model.add(BatchNormalization())
self.model.add(Dropout(rate=dropout_rate))
self.model.add(Dense(n_classes, activation='softmax'))
from tensorflow.keras.activations import softplus, relu
from tensorflow.keras.backend import expand_dims, clip
from tensorflow.keras.layers import Input, Dense, Concatenate
from tensorflow.keras.losses import MeanSquaredError
from tensorflow.keras.models import Model
from tensorflow_probability.python.distributions import Normal
env = gym.make('MountainCarContinuous-v0')
state = env.reset()
state_input = Input(state.shape, name='state_input')
actor_dense1_layer = Dense(4, activation=relu, name='actor_dense1')
actor_dense1 = actor_dense1_layer(state_input)
actor_out_layer = Dense(2, activation=None, name='actor_nn_out')
actor_out = actor_out_layer(actor_dense1)
mean = actor_out[:, 0]
std = softplus(actor_out[:, 1])
dist = Normal(mean, std, name='dist')
action = dist.sample((), name='sample')
action_log_prob = dist.log_prob(action, name='log_prob')
action_log_prob = expand_dims(action_log_prob)
action = expand_dims(action)
action = clip(action, -1.0, 1.0)
critic_concat = Concatenate()([state_input, action])
critic_dense1_layer = Dense(units=8, activation=relu, name='critic_dense1')
critic_dense1 = critic_dense1_layer(critic_concat)
critic_out_layer = Dense(1, activation=None)
critic_out = critic_out_layer(critic_dense1)
actor_critic_model = Model(inputs=[state_input],
outputs=[action, action_log_prob, critic_out])
critic_loss_op = MeanSquaredError(name='critic_loss')
actor_vars = actor_dense1_layer.trainable_variables + actor_out_layer.trainable_variables
def _training_loop(
self,
filepath: str,
train_gen: train_ts_generator, # can name of function be type?
val_gen: train_ts_generator,
epochs: int = 100,
steps_per_epoch: int = 50,
early_stopping: int = True,
stopping_patience: int = 5,
stopping_delta: int = 1,
) -> typing.Tuple[tf.Tensor, int]:
"""
util function
iterates over batches, updates gradients, records metrics, writes to tb, checkpoints, early stopping
"""
# set up metrics to track during training
batch_loss_avg = Mean()
epoch_loss_avg = Mean()
eval_loss_avg = Mean()
eval_mae = MeanAbsoluteError()
eval_rmse = RootMeanSquaredError()
# set up early stopping callback
early_stopping_cb = EarlyStopping(patience=stopping_patience,
active=early_stopping,
delta=stopping_delta)
# setup table for unscaling
self._lookup_table = build_tf_lookup(self._ts_obj.target_means)
# Iterate over epochs.
best_metric = math.inf
for epoch in range(epochs):
logger.info(f"Start of epoch {epoch}")
start_time = time.time()
for batch, (x_batch_train, cat_labels,
y_batch_train) in enumerate(train_gen):
# compute loss
with tf.GradientTape(persistent=True) as tape:
mu, scale = self._model(x_batch_train, training=True)
# softplus parameters
scale = softplus(scale)
if self._ts_obj.count_data:
mu = softplus(mu)
mu, scale = unscale(mu, scale, cat_labels,
self._lookup_table)
loss_value = self._loss_fn(y_batch_train, (mu, scale))
# sgd
if self._tb:
tf.summary.scalar("train_loss", loss_value,
epoch * steps_per_epoch + batch)
batch_loss_avg(loss_value)
epoch_loss_avg(loss_value)
grads = tape.gradient(loss_value,
self._model.trainable_weights)
self._optimizer.apply_gradients(
zip(grads, self._model.trainable_weights))
# Log 5x per epoch.
if batch % (steps_per_epoch // 5) == 0 and batch != 0:
logger.info(
f"Epoch {epoch}: Avg train loss over last {(steps_per_epoch // 5)} steps: {batch_loss_avg.result()}"
)
batch_loss_avg.reset_states()
# Run each epoch batches times
epoch_loss_avg_result = epoch_loss_avg.result()
if batch == steps_per_epoch:
logger.info(
f"Epoch {epoch} took {round(time.time() - start_time, 0)}s : Avg train loss: {epoch_loss_avg_result}"
)
break
# validation
if val_gen is not None:
logger.info(f"End of epoch {epoch}, validating...")
start_time = time.time()
for batch, (x_batch_val, cat_labels,
y_batch_val) in enumerate(val_gen):
# compute loss, doesn't need to be persistent bc not updating weights
with tf.GradientTape() as tape:
# treat as training -> reset lstm states inbetween each batch
mu, scale = self._model(x_batch_val, training=True)
# softplus parameters
mu, scale = self._softplus(mu, scale)
# unscale parameters
mu, scale = unscale(mu, scale, cat_labels,
self._lookup_table)
# calculate loss
loss_value = self._loss_fn(y_batch_val, (mu, scale))
# log validation metrics (avg loss, avg MAE, avg RMSE)
eval_mae(y_batch_val, mu)
eval_rmse(y_batch_val, mu)
eval_loss_avg(loss_value)
if batch == steps_per_epoch:
break
# logging
eval_mae_result = eval_mae.result()
logger.info(
f"Validation took {round(time.time() - start_time, 0)}s")
logger.info(
f"Epoch {epoch}: Val loss on {steps_per_epoch} steps: {eval_loss_avg.result()}"
)
logger.info(
f"Epoch {epoch}: Val MAE: {eval_mae_result}, RMSE: {eval_rmse.result()}"
)
if self._tb:
tf.summary.scalar("val_loss", eval_loss_avg.result(),
epoch)
tf.summary.scalar("val_mae", eval_mae_result, epoch)
tf.summary.scalar("val_rmse", eval_rmse.result(), epoch)
new_metric = eval_mae_result
# early stopping
if early_stopping_cb(eval_mae_result):
break
# reset metric states
eval_loss_avg.reset_states()
eval_mae.reset_states()
eval_rmse.reset_states()
else:
if early_stopping_cb(epoch_loss_avg_result):
break
new_metric = epoch_loss_avg_result
# update best_metric and save new checkpoint if improvement
if new_metric < best_metric:
best_metric = new_metric
if filepath is not None:
self._checkpointer.save(file_prefix=filepath)
else:
self.save_weights("model_best_weights.h5")
# reset epoch loss metric
epoch_loss_avg.reset_states()
# load in best weights before returning if not using checkpointer
if filepath is None:
self.load_weights("model_best_weights.h5")
os.remove("model_best_weights.h5")
return best_metric, epoch + 1