for batch, (train_x, train_y) in enumerate(train_dataset):
# print(train_x.shape, train_y.shape)
# assert train_x.shape == (batch_size, n_steps, 1)
# assert train_y.shape == (batch_size, n_steps)
# loss
with tf.GradientTape() as tape:
y = model(train_x, training=True)
# assert y.shape == (batch_size, n_steps, 10)
loss_np = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=train_y, logits=y)
# assert loss_np.shape == (batch_size, n_steps)
loss = tf.reduce_mean(loss_np)
# gradient
gradient = tape.gradient(loss, model.trainable_variables)
if clip != -1:
gradient, _ = tf.clip_by_global_norm(gradient, clip)
optimizer.apply_gradients(zip(gradient, model.trainable_variables))
if batch % 100 == 0:
print("Batch:", batch, ", Train loss:", loss.numpy())
# Eval
eval_logits = model(test_data, training=False)
eval_loss_list = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=test_labels, logits=eval_logits)
# assert eval_loss_list.shape == (batch_size, n_steps)
eval_loss = tf.reduce_mean(eval_loss_list)
print("Epoch:", epoch, ", Eval loss:", eval_loss.numpy(), "\n---\n")
# Save
model.save_weights("weights/model_weight.h5")
gradient = tape.gradient(loss, model.trainable_variables)
if clip != -1:
gradient, _ = tf.clip_by_global_norm(gradient, clip)
optimizer.apply_gradients(zip(gradient, model.trainable_variables))
if batch % 100 == 0:
print("Batch:", batch, ", Train loss:", loss.numpy())
# Eval Acc
eval_labels = model(test_data, training=False)
eval_acc = np.mean(
np.argmax(eval_labels.numpy(), axis=1) == test_labels.numpy())
print("Epoch:", epoch, ", Eval acc:", eval_acc * 100, "%\n---\n")
# Save
if not args.permute:
model.save_weights("Sequential-weights/model_weight.h5")
else:
model.save_weights("Permuted-weights/model_weight.h5")
# Eval Acc
# print(test_data.shape, test_labels.shape)
# test_data = test_data[:100, :, :]
# test_labels = test_labels[:100]
# print("Before:")
# eval_labels = model(test_data, training=False)
# eval_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=test_labels, logits=eval_labels))
# print("eval_loss:", eval_loss)
# print("\nAfter:")