I am not using model.compile() function, and I am following 'Writing a training loop from scratch' (Therefore, simple 'model.save()' does not include my optimizer). Also, I create NN using a functional API in Tensorflow. So, I am able to save using 'model.save()', but it saves only NN model and does not save the state of my optimizer.
I think I was able to save an optimizer using pickle, but somehow I am no longer able to save it before
(it give me an error message (AttributeError: Can't pickle local object 'make_gradient_clipnorm_fn..'))
I can use 'get_config' and 'from_config' methods, but it does not save weights. There is another solution to save weights like here. I am looking for a better (and simple) solution.
(and actually this method does not work for me because I am not only training 'model.trainable_weights', but also other 'tf.Variable', which are not belong to my 'model'.)
Is there any easy and convenient way to save and load optimizer?
Thanks
colab: https://colab.research.google.com/drive/1Jn2lbMcaVURpKITL6qEQ-05XndVDoABQ?usp=sharing
or --code--
import numpy as np
import tensorflow as tf
import matplotlib
import matplotlib.pyplot as plt
import datetime
import os
import importlib
import pickle
import time
from pathlib import Path
gpus = tf.config.experimental.list_physical_devices('GPU')
print(gpus)
"""Generate samples for training"""
my_dtype = 'float32'
tf.keras.backend.set_floatx(my_dtype)
num_smpl = 500
amp1 = 0.5
x = 0.5*np.random.rand(num_smpl, 1).astype(np.float32)
y = amp1 * np.heaviside(x, 0) - amp1*0.2*np.heaviside(x-0.15, 0) \
+ amp1*0.2*np.heaviside(x-0.17, 0)
mu, sigma = 0, 0.001 # mean and standard deviation
white_noise = np.random.normal(mu, sigma, num_smpl).reshape(-1, 1)
y = y + white_noise
fig = plt.figure()
ax = fig.add_subplot(111)
ax.scatter(x, y, s=5, edgecolors='k')
"""Create model using a funtional API"""
layer_input = tf.keras.layers.Input(shape=(1,), dtype=my_dtype)
layer_dense = tf.keras.layers.Dense(8, activation='tanh', dtype=my_dtype)(layer_input)
layer_dense = tf.keras.layers.Dense(8, activation='tanh', dtype=my_dtype)(layer_dense)
layer_output = tf.keras.layers.Dense(1, name='', dtype=my_dtype)(layer_dense)
model = tf.keras.Model(inputs=layer_input, outputs=layer_output, name='my_model')
model.summary()
epoch = 1
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01, epsilon=1e-09)
loss_fn = tf.keras.losses.MeanSquaredError()
"""Training using a custom training loop"""
while epoch < 100:
with tf.GradientTape() as t1:
output = model(x, training=True)
mse = loss_fn(output, y)
grads = t1.gradient(mse, model.trainable_weights)
optimizer.apply_gradients(zip(grads, model.trainable_weights))
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
if epoch % 10 == 0:
print('Epoch: {}, Loss: {}'.format(epoch, mse))
epoch = epoch + 1
"""Normal variables can be saved using pickle"""
print('Cur. locaiton: ' + os.getcwd())
a = np.arange(0., 10.)
print(a)
with open('test.pkl', 'wb') as file:
pickle.dump(a, file)
"""***How to save optimizer?***"""
with open('my_opti.pkl', 'wb') as file:
pickle.dump(optimizer, file)
'
