def compile(self,
model_config: dict,
optimizer_config: dict,
max_to_keep: int = 10):
with self.strategy.scope():
self.generator = create_generator(
g_enc_depths=model_config["g_enc_depths"],
window_size=self.speech_config["window_size"],
kwidth=model_config["kwidth"],
ratio=model_config["ratio"])
self.discriminator = create_discriminator(
d_num_fmaps=model_config["d_num_fmaps"],
window_size=self.speech_config["window_size"],
kwidth=model_config["kwidth"],
ratio=model_config["ratio"])
self.generator.summary(line_length=150)
self.generator_optimizer = tf.keras.optimizers.get(
optimizer_config["generator"])
self.discriminator_optimizer = tf.keras.optimizers.get(
optimizer_config["discriminator"])
if self.is_mixed_precision:
self.generator_optimizer = mixed_precision.LossScaleOptimizer(
self.generator_optimizer, "dynamic")
self.discriminator_optimizer = mixed_precision.LossScaleOptimizer(
self.discriminator_optimizer, "dynamic")
self.create_checkpoint_manager(
max_to_keep,
generator=self.generator,
gen_optimizer=self.generator_optimizer,
discriminator=self.discriminator,
disc_optimizer=self.discriminator_optimizer)
def _instantiate_optimizers(strategy, learning_rate, beta_1, train_settings):
LOGGER.info(" -------- Creating Optimizers --------")
with strategy.scope():
srfr_optimizer = NovoGrad(
learning_rate=learning_rate,
beta_1=beta_1,
beta_2=train_settings["beta_2"],
weight_decay=train_settings["weight_decay"],
name="novograd_srfr",
)
srfr_optimizer = mixed_precision.LossScaleOptimizer(
srfr_optimizer,
loss_scale="dynamic",
)
discriminator_optimizer = NovoGrad(
learning_rate=learning_rate,
beta_1=beta_1,
beta_2=train_settings["beta_2"],
weight_decay=train_settings["weight_decay"],
name="novograd_discriminator",
)
discriminator_optimizer = mixed_precision.LossScaleOptimizer(
discriminator_optimizer, loss_scale="dynamic")
return (
srfr_optimizer,
discriminator_optimizer,
)
def _instantiate_optimizers(self, learning_rate, beta_1, train_settings):
self.logger.info(" -------- Creating Optimizers --------")
with self.strategy.scope():
srfr_optimizer = AdamW(
learning_rate=learning_rate,
beta_1=beta_1,
beta_2=train_settings["beta_2"],
weight_decay=train_settings["weight_decay"],
name="adam_srfr",
)
srfr_optimizer = mixed_precision.LossScaleOptimizer(
srfr_optimizer,
loss_scale="dynamic",
)
discriminator_optimizer = AdamW(
learning_rate=learning_rate,
beta_1=beta_1,
beta_2=train_settings["beta_2"],
weight_decay=train_settings["weight_decay"],
name="adam_discriminator",
)
discriminator_optimizer = mixed_precision.LossScaleOptimizer(
discriminator_optimizer, loss_scale="dynamic")
return (
srfr_optimizer,
discriminator_optimizer,
)
def __init__(self,
encoder,
decoder,
config):
self.encoder= encoder
self.decoder= decoder
self.config= config
self.ckpt_dir= config['rnn_attention']['ckpt_dir']
self.dec_max_len= config['dataloader']['dec_max_len']
self.epochs= config['rnn_attention']['epochs']
self.learning_rate= config['rnn_attention']['learning_rate']
self.optimizer= OPTM(config['rnn_attention']['optimizer'], self.learning_rate)
self.optimizer= mixed_precision.LossScaleOptimizer(self.optimizer, loss_scale= 'dynamic')
self.sce= tf.keras.losses.SparseCategoricalCrossentropy(from_logits= True, reduction= 'none')
ckpt_dir= config['rnn_attention']['ckpt_dir']
precision_policy= config['rnn_attention']['policy']
super(TrainerRNNAttention, self).__init__(ckpt_dir,
precision_policy,
model1= encoder,
model2= decoder,
optimizer= self.optimizer,
name= 'rnn_attention')
def __init__(self,
transformer,
config):
self.transformer= transformer
self.config= config
self.batch_size= self.config['dataloader']['batch_size']
self.dmodel= config['transformer']['dmodel']
self.num_heads= config['transformer']['num_heads']
self.depth= self.dmodel/ self.num_heads
self.warmup_steps= config['transformer']['warmup_steps']
if config['transformer']['learning_rate'] == 'schedule':
self.learning_rate= LearningRate(self.dmodel, self.warmup_steps)
else:
self.learning_rate= config['transformer']['learning_rate']
self.optimizer= OPTM(config['transformer']['optimizer'], self.learning_rate)
self.optimizer= mixed_precision.LossScaleOptimizer(self.optimizer, loss_scale= 'dynamic')
self.sce= tf.keras.losses.SparseCategoricalCrossentropy(from_logits= True, reduction= 'none')
ckpt_dir= config['transformers']['ckpt_dir']
precision_policy= config['transformers']['policy']
super(TrainerTransformer, self).__init__(ckpt_dir,
precision_policy,
model1= transformer,
optimizer= self.optimizer
name= 'transformer')
def multi_classifier_training(s, autoencoder, classifiers, train_dataset,
val_dataset, test_dataset):
anomaly_detector = PerceptualMultiClassifier(autoencoder, *classifiers)
loss_function = define_loss_function(s.LOSS_FUNCTION)
optimizer = define_optimizer(s.OPTIMIZER, s.OPTIMIZER_PARAMS)
optimizer = mixed_precision.LossScaleOptimizer(optimizer,
loss_scale='dynamic')
anomaly_detector.compile(optimizer=optimizer, loss=loss_function)
model_save_time = datetime.today().strftime(
'%Y-%m-%d') + '_' + datetime.now().time().strftime('%H-%M-%S')
model_path = s.PROJECT_MODEL_PATH + model_save_time
inference_cb = CustomInference2(val_data=val_dataset,
save_model_path=model_path,
show_hist=s.SHOW_HISTOGRAMS_PLOTS)
anomaly_detector.fit(train_dataset,
validation_data=val_dataset,
epochs=s.EPOCHS,
callbacks=[inference_cb])
# restore weights
best_ae_weights, best_clf1_weights, best_clf2_weights, _ = extract_best_weights(
model_path)
anomaly_detector.autoencoder.load_weights(best_ae_weights)
anomaly_detector.classifier1.load_weights(best_clf1_weights)
anomaly_detector.classifier2.load_weights(best_clf2_weights)
cb = CustomInference2(val_data=test_dataset,
save_model_path=None,
show_hist=s.SHOW_HISTOGRAMS_PLOTS)
anomaly_detector.evaluate(test_dataset, callbacks=[cb])
def init(self):
self.model = HResResNet18(preact=True)
train_metrics = {
'loss':
tf.keras.metrics.MeanTensor('loss'),
'rmse':
metrics.RootMeanSquaredError('rmse', denorm_fn=self.y_denorm_fn),
'mae':
metrics.MeanAbsoluteError('mae', denorm_fn=self.y_denorm_fn)
}
valid_metrics = {
'loss':
tf.keras.metrics.MeanTensor('loss'),
'rmse':
metrics.RootMeanSquaredError('rmse', denorm_fn=self.y_denorm_fn),
'mae':
metrics.MeanAbsoluteError('mae', denorm_fn=self.y_denorm_fn)
}
self.optim = tfa.optimizers.AdamW(weight_decay=FLAGS.wd,
lr=0.0,
beta_1=FLAGS.adam_beta_1,
beta_2=FLAGS.adam_beta_2,
epsilon=FLAGS.adam_epsilon)
if FLAGS.amp:
self.optim = mixed_precision.LossScaleOptimizer(
self.optim, loss_scale='dynamic')
return {'HResResNet18': self.model}, None, train_metrics, valid_metrics
def _create_optimizer(self):
"""Creates optimizer."""
params = self.params
# TODO(b/139414679): Explore the difference between using
opt = tf.keras.optimizers.Adam(
params["learning_rate"],
params["optimizer_adam_beta1"],
params["optimizer_adam_beta2"],
epsilon=params["optimizer_adam_epsilon"])
if params["dtype"] == tf.float16:
opt = mixed_precision.LossScaleOptimizer(
opt,
loss_scale=flags_core.get_loss_scale(
self.flags_obj, default_for_fp16="dynamic"))
if self.flags_obj.fp16_implementation == "graph_rewrite":
# Note: when flags_obj.fp16_implementation == "graph_rewrite", dtype as
# determined by flags_core.get_tf_dtype(flags_obj) would be 'float32'
# which will ensure tf.compat.v2.keras.mixed_precision and
# tf.train.experimental.enable_mixed_precision_graph_rewrite do not double
# up.
opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(
opt)
return opt
def init(self):
self.model = ResNet18(feature_mode=True, preact=True, last_norm=True)
self.tail = ResNetTail(preact=True, last_norm=True)
train_metrics = {
'loss': tf.keras.metrics.Mean('loss'),
'acc': tf.keras.metrics.SparseCategoricalAccuracy('acc')
}
valid_metrics = {
'loss': tf.keras.metrics.Mean('loss'),
'acc': tf.keras.metrics.SparseCategoricalAccuracy('acc')
}
self.loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True,
reduction=tf.keras.losses.Reduction.NONE) # distributed-aware
self.optim = tfa.optimizers.AdamW(weight_decay=FLAGS.wd,
lr=0.0,
beta_1=FLAGS.adam_beta_1,
beta_2=FLAGS.adam_beta_2,
epsilon=FLAGS.adam_epsilon)
if FLAGS.amp:
self.optim = mixed_precision.LossScaleOptimizer(
self.optim, loss_scale='dynamic')
return {
'pretrained_resnet18': self.model,
'tail': self.tail
}, {
'pretrained_resnet18': [32, 32, 3],
'tail': [4, 4, 512]
}, train_metrics, valid_metrics
def __init__(self, hparams):
self._mixed_precision = hparams.mixed_precision
optimizer = tf.keras.optimizers.Adam(hparams.learning_rate)
if self._mixed_precision:
optimizer = mixed_precision.LossScaleOptimizer(
optimizer, loss_scale='dynamic')
self.optimizer = optimizer
def init_optimizer(self):
self.init_learning_rate()
if 'RMS' in self.args.optimizer.upper():
# Use RMS prop:
self.print("Selected optimizer is RMS Prop")
self._opt = tf.compat.v1.train.RMSPropOptimizer(
self._learning_rate)
else:
# default is Adam:
self.print("Using default Adam optimizer")
self._opt = tf.compat.v1.train.AdamOptimizer(self._learning_rate)
if self.args.precision == "mixed":
from tensorflow.keras.mixed_precision import experimental as mixed_precision
self._opt = mixed_precision.LossScaleOptimizer(
self._opt, loss_scale='dynamic')
else:
with tf.name_scope('gradient_accumulation'):
self._zero_gradients = [
tv.assign(tf.zeros_like(tv)) for tv in self._accum_vars
]
self._accum_gradients = [
self._accum_vars[i].assign_add(gv[0])
for i, gv in enumerate(
self._opt.compute_gradients(self._loss))
]
self._apply_gradients = self._opt.apply_gradients(
zip(self._accum_vars, tf.compat.v1.trainable_variables()),
global_step=self._global_step)
def __init__(self,
model,
criterion,
optimizers,
train_metrics: Mapping[str, Metric],
eval_metrics: Mapping[str, Metric],
work_dir,
grad_clip_norm=0.0,
multiple_steps=True,
xla_compile=False,
output_transform=default_metric_transform,
n_batches_per_step=None):
if not isinstance(optimizers, Sequence):
optimizers = [optimizers]
optimizers = list(optimizers)
self._strategy = parse_strategy('auto')
work_dir = fmt_path(work_dir)
self.model = model
self.criterion = criterion
self.optimizers = optimizers
self.train_metrics = train_metrics
self.eval_metrics = eval_metrics
self.work_dir = work_dir
self.dtype = tf.dtypes.as_dtype(
mixed_precision.global_policy().compute_dtype)
if self.dtype == tf.float16:
self.optimizers = [
mixed_precision.LossScaleOptimizer(optimizer, 'dynamic') if
not isinstance(optimizer, mixed_precision.LossScaleOptimizer)
else optimizer for optimizer in self.optimizers
]
self.grad_clip_norm = grad_clip_norm
self.multiple_steps = multiple_steps
self.xla_compile = xla_compile
self.output_transform = output_transform
self._log_dir = self.work_dir / "runs"
self._writer = None
self._verbose = True
self._state = {
"train": {},
"eval": {},
"test": {},
}
# epoch -> stage -> metric -> value
self.metric_history = MetricHistory(["train", "eval", "test"])
self._terminated = False
self.set_global_state("epoch", tf.Variable(-1, dtype=tf.int64))
if self.xla_compile:
self.xla_train_batch = tf.function(self.train_batch,
experimental_compile=True)
self.n_batches_per_step = n_batches_per_step
def TransformerOptimizer(initial_lr, **kwargs):
adam_opt = tf.keras.optimizers.Adam(CustomSchedule(initial_lr, **kwargs),
beta_1=0.9,
beta_2=0.98,
epsilon=1e-9)
return mixed_precision.LossScaleOptimizer(adam_opt, loss_scale='dynamic')
def init_optimizer(self):
"""
Init self.optimizer.
"""
self.learning_rate = tf.keras.optimizers.schedules.ExponentialDecay(self.learning_rate, self.max_iter, 0.1)
self.optimizer = tf.keras.optimizers.Adam(learning_rate=self.learning_rate)
if self.use_mixed_precision:
self.optimizer = mixed_precision.LossScaleOptimizer(self.optimizer,
loss_scale=tf.mixed_precision.experimental.DynamicLossScale(initial_loss_scale=2 ** 15, increment_period=1000))
def __init__(self, name, modules, lr, clip=None, wd=None, wdpattern=r'.*'):
self._name = name
self._modules = modules
self._clip = clip
self._wd = wd
self._wdpattern = wdpattern
self._opt = tf.optimizers.Adam(lr)
self._opt = prec.LossScaleOptimizer(self._opt, 'dynamic')
self._variables = None
def __init__(self,
action_size,
obvs_size,
latent_size=512,
mem_size=10000):
super(SAC, self).__init__()
self.action_size = action_size
self.latent_size = latent_size
self.lr = 0.0005
self.counter = 1
self.delay = 2
self.batchsize = 32
self.tau = 0.005
self.gamma = 0.99
self.grad_clip = 0.5
self.target_ent = float(-action_size)
self.target_ent = tf.Variable(self.target_ent, trainable=False)
self.alpha = tf.Variable(0.2, trainable=False)
self.log_alpha = tf.Variable(tf.math.log(self.alpha), trainable=True)
self.buffer = {
'states': np.zeros((mem_size, ) + obvs_size, dtype=np.float32),
'actions': np.zeros((mem_size, action_size), dtype=np.float32),
'rewards': np.zeros((mem_size, ), dtype=np.float32),
'next_states': np.zeros((mem_size, ) + obvs_size,
dtype=np.float32),
'dones': np.zeros((mem_size, ), dtype=np.float32),
}
self.replay_buffer = Replay_Buffer(self.buffer, capacity=mem_size)
self.agent = Agent(action_size=self.action_size)
self.critic1 = Critic(action_size=self.action_size)
self.critic2 = Critic(action_size=self.action_size)
self.target_agent = Agent(action_size=self.action_size)
self.target_critic1 = Critic(action_size=self.action_size)
self.target_critic2 = Critic(action_size=self.action_size)
self.agent_optimizer = mixed_precision.LossScaleOptimizer(
Adam(self.lr, epsilon=1e-8), loss_scale='dynamic')
self.critic1_optimizer = mixed_precision.LossScaleOptimizer(
Adam(self.lr, epsilon=1e-8), loss_scale='dynamic')
self.critic2_optimizer = mixed_precision.LossScaleOptimizer(
Adam(self.lr, epsilon=1e-8), loss_scale='dynamic')
self.alpha_optimizer = mixed_precision.LossScaleOptimizer(
Adam(self.lr, epsilon=1e-8), loss_scale='dynamic')
def compile(self, optimizer, loss_fn, metrics, **kwargs):
if USE_FLOAT16:
optimizer = mixed_precision.LossScaleOptimizer(
optimizer, loss_scale='dynamic')
super(Captioner, self).compile(optimizer=optimizer,
loss=loss_fn,
**kwargs)
# self.optimizer = optimizer
# self.loss_fn = loss_fn
self.word_metrics = metrics
def compile(self,
model: Transducer,
optimizer: any,
max_to_keep: int = 10):
with self.strategy.scope():
self.model = model
self.model.summary(line_length=100)
self.optimizer = tf.keras.optimizers.get(optimizer)
if self.is_mixed_precision:
self.optimizer = mixed_precision.LossScaleOptimizer(
self.optimizer, "dynamic")
self.create_checkpoint_manager(max_to_keep,
model=self.model,
optimizer=self.optimizer)
def compile(self, model_g, opt_g, model_d=None, opt_d=None):
with self.strategy.scope():
self.model_g = model_g
self.model_g._build()
self.mels_loss = TFMultiResolutionSTFT()
self.model_g.summary(line_length=100)
self.optimizer_g = tf.keras.optimizers.get(opt_g)
if self.config['use_gan']:
self.model_d = model_d
self.model_d._build()
self.optimizer_d = tf.keras.optimizers.get(opt_d)
if self.is_mixed_precision:
self.optimizer_g = mixed_precision.LossScaleOptimizer(
self.optimizer_g, "dynamic")
if self.config['use_gan']:
self.optimizer_d = mixed_precision.LossScaleOptimizer(
self.optimizer_d, "dynamic")
try:
self.load_checkpoint()
except:
logging.info('trainer resume failed')
self.set_progbar()
def compile(self, model, opt):
with self.strategy.scope():
self.model = model
self.model._build()
try:
self.load_checkpoint()
except:
logging.info('trainer resume failed')
self.model.summary(line_length=100)
self.optimizer = tf.keras.optimizers.get(opt)
if self.is_mixed_precision:
self.optimizer = mixed_precision.LossScaleOptimizer(self.optimizer, "dynamic")
self.set_progbar()
def loss_scale_optimizer(
cls, optimizer: keras.optimizers.Optimizer
) -> mixedprecision.LossScaleOptimizer:
""" Optimize loss scaling for mixed precision training.
Parameters
----------
optimizer: :class:`tf.keras.optimizers.Optimizer`
The optimizer instance to wrap
Returns
--------
:class:`tf.keras.mixed_precision.loss_scale_optimizer.LossScaleOptimizer`
The original optimizer with loss scaling applied
"""
return mixedprecision.LossScaleOptimizer(optimizer)
def _optimizer_fn_to_optimizer(
optimizer_fn: Union[Callable, None], model: Model, framework: str
) -> Union[None, tf.optimizers.Optimizer, torch.optim.Optimizer]:
"""A helper function to invoke an optimizer function.
Args:
optimizer_fn: The function to be invoked in order to instantiate an optimizer.
model: The model with which the optimizer should be associated.
framework: Which backend framework should be used ('tf' or 'torch').
Returns:
An optimizer instance, or None if `optimizer_fn` was None.
"""
optimizer = None
if optimizer_fn:
if framework == "tf":
try:
optimizer = optimizer_fn()
except:
raise AssertionError(
"optimizer_fn of Tensorflow backend should be callable without args. Please "
"make sure model and optimizer_fn are using the same backend"
)
# initialize optimizer variables
_ = optimizer.iterations
optimizer._create_hypers()
optimizer._create_slots(model.trainable_variables)
# handle mixed precision loss scaling
if mixed_precision_tf.global_policy().name != "float32":
optimizer = mixed_precision_tf.LossScaleOptimizer(
optimizer, loss_scale='dynamic')
assert isinstance(
optimizer, tf.optimizers.Optimizer
), "optimizer_fn should generate tensorflow optimizer"
else:
try:
optimizer = optimizer_fn(model.parameters())
except Exception as e:
print(
"optimizer_fn of Pytorch backend should be callable with single arg. Please sure model and \
optimizer_fn are using the same backend")
raise ValueError(repr(e))
assert isinstance(
optimizer, torch.optim.Optimizer
), "optimizer_fn should generate pytorch optimizer"
return optimizer
def build_anomaly_detection_model(s,
auto_encoder,
classifier,
show_graph=False):
anomaly_detector = PerceptualAutoEncoder(auto_encoder, classifier)
if show_graph:
auto_encoder.summary()
classifier.summary()
optimizer = define_optimizer(s.OPTIMIZER, s.OPTIMIZER_PARAMS)
optimizer = mixed_precision.LossScaleOptimizer(optimizer,
loss_scale='dynamic')
loss_function = define_loss_function(s.LOSS_FUNCTION)
anomaly_detector.compile(optimizer=optimizer, loss=loss_function)
return anomaly_detector
def build_optimizer(step=None):
schedule = get_lr_schedule()
optimizer = get_optimizer()
if FLAGS.optimizer in ['adamw']:
if FLAGS.schedule == 'constant':
opt = optimizer(learning_rate=schedule,
weight_decay=FLAGS.weight_decay)
else:
raise NotImplementedError
else:
opt = optimizer(learning_rate=schedule)
if FLAGS.use_lookahead:
opt = tfa.optimizers.Lookahead(opt, sync_period=6, slow_step_size=0.5)
if FLAGS.mixed_precision:
opt = mixed_precision.LossScaleOptimizer(opt, loss_scale='dynamic')
return opt
def compile(self,
model: Transducer,
optimizer: any,
max_to_keep: int = 10):
f, c = self.speech_featurizer.compute_feature_dim()
with self.strategy.scope():
self.model = model
self.model._build([1,80,f,c])
self.model.summary(line_length=100)
try:
self.load_checkpoint()
except:
logging.info('trainer resume failed')
self.optimizer = tf.keras.optimizers.get(optimizer)
if self.is_mixed_precision:
self.optimizer = mixed_precision.LossScaleOptimizer(self.optimizer, "dynamic")
self.set_progbar()
def init_optimizer(self):
self.init_learning_rate()
if self.args.mode.optimizer.name == "rmsprop":
# Use RMS prop:
self._opt = tf.keras.optimizers.RMSprop(self._learning_rate)
else:
# default is Adam:
self._opt = tf.keras.optimizers.Adam(self._learning_rate)
if self.args.run.precision == "mixed":
from tensorflow.keras.mixed_precision import experimental as mixed_precision
self._opt = mixed_precision.LossScaleOptimizer(
self._opt, loss_scale='dynamic')
self.tape = tf.GradientTape()
def init(self):
valid_amount = 0
for x_batch, _ in self.valid_dataset:
valid_amount = valid_amount + x_batch.shape[0]
input_shape = self.train_dataset.element_spec[0][0].shape[
1:] # [0][0] -> ((x, x_blur),y) -> x
self.predictor = Predictor()
self.randnet_1 = RandomNet() # Note: fixed, not trainable
self.randnet_2 = RandomNet() # Note: fixed, not trainable
self.norm_all = layers.Norm(center=False, scale=False)
self.randnet_1.trainable = False
self.randnet_2.trainable = False
train_metrics = {
'loss': tf.keras.metrics.Mean('loss'),
}
valid_metrics = {
'figure':
metrics.RNDMetrics(mode=metrics.RNDMetrics.MODE_FIGURE,
amount=valid_amount),
'tnr@95tpr':
metrics.RNDMetrics(mode=metrics.RNDMetrics.MODE_TNR95TPR,
amount=valid_amount)
}
self.loss_object = losses.MSELoss()
self.optim = tfa.optimizers.AdamW(weight_decay=FLAGS.wd,
lr=0.0,
beta_1=FLAGS.adam_beta_1,
beta_2=FLAGS.adam_beta_2,
epsilon=FLAGS.adam_epsilon)
if FLAGS.amp:
self.optim = mixed_precision.LossScaleOptimizer(
self.optim,
loss_scale=tf.mixed_precision.experimental.DynamicLossScale(
initial_loss_scale=(2**15),
increment_period=20,
multiplier=2.0))
return {
'predictor': self.predictor,
'randnet_1': self.randnet_1,
'randnet_2': self.randnet_2
}, {
'predictor': input_shape,
'randnet_1': input_shape,
'randnet_2': input_shape
}, train_metrics, valid_metrics
def compile(self, model: tf.keras.Model,
optimizer: any,
max_to_keep: int = 10):
with self.strategy.scope():
self.model = model
self.optimizer = tf.keras.optimizers.get(optimizer)
self.model._build()
try:
self.load_checkpoint()
except:
logging.info('lm trainer resume failed')
if self.is_mixed_precision:
self.optimizer = mixed_precision.LossScaleOptimizer(self.optimizer, "dynamic")
self.set_progbar()
def get_compiled_model(hyper_parameters):
model_save_dir = hyper_parameters["model_save_dir"]
model_name = hyper_parameters["model_name"]
metrics = get_metrics()
if hyper_parameters["saved_model_path"] is not None:
print("Continuing training model from last checkpoint")
model = load_trained_tf_model(
model_save_dir=hyper_parameters["model_save_dir"],
model_name=hyper_parameters["saved_model_path"],
use_best_val_loss=True,
)
else:
# Build the model.
# TODO try using this for training very large models:
# https://keras.io/getting-started/faq/#device-parallelism
model = hyper_parameters["model"](
input_shape=hyper_parameters["input_shape"],
print_model_summary=hyper_parameters["print_model_summary"],
)
# Define the optimiser.
optimiser = hyper_parameters["optimiser"](
lr=hyper_parameters["learning_rate"]
)
# Enable mixed precision using dynamic loss scaling.
if hyper_parameters["enable_amp"]:
optimiser = mixed_precision.LossScaleOptimizer(
optimiser, loss_scale="dynamic"
)
# Compile the model.
model.compile(
optimizer=optimiser, loss="categorical_crossentropy", metrics=metrics,
)
# Serialize the model to JSON and save.
model_json = model.to_json()
with open(f"{model_save_dir}{model_name}.json", "w") as json_file:
json_file.write(model_json)
return model
def init_optimizer(self):
self.init_learning_rate()
if 'RMS' in self.args.optimizer.upper():
# Use RMS prop:
self.print("Selected optimizer is RMS Prop")
self._opt = tf.keras.optimizers.RMSprop(self._learning_rate)
else:
# default is Adam:
self.print("Using default Adam optimizer")
self._opt = tf.keras.optimizers.Adam(self._learning_rate)
if self.args.precision == "mixed":
from tensorflow.keras.mixed_precision import experimental as mixed_precision
self._opt = mixed_precision.LossScaleOptimizer(
self._opt, loss_scale='dynamic')
self.tape = tf.GradientTape()
