def main():
config = argparse.parse_cmd(arguments)
datasets = {
'train':
dataloader.get_dataset(
config['dataloader'],
transformation_list=config['dataloader']['train_list'],
num_classes=10,
split=config['dataloader']['train_split_id']),
'val':
dataloader.get_dataset(
config['dataloader'],
transformation_list=config['dataloader']['val_list'],
num_classes=10,
split=config['dataloader']['val_split_id'])
}
for dataset_type in ['train', 'val']:
for i, (images, y) in enumerate(datasets[dataset_type]):
image = images[0]
# opencv manage images as BGR object, TF as RGB
image = image.numpy()[:, :, ::-1]
cv2.imwrite(os.path.join('/tmp', f'{dataset_type}_{i}.png'), image)
if i == 20:
break
def test_compare_without_augmentations(self):
# create fake dataset and convert TFrecords
dataset_dir = create_fake_dataset_and_convert_to_tfrecords(
image_array=self.image_white, n_images_per_class=10)
config = {
'name': 'tfrecords',
'data_dir': dataset_dir,
'batch_size': BATCH_SIZE,
'train_split_id': 'train',
}
# get data from tfrecord
dataset = get_dataset(config, [], 1, 'train')
image, _ = next(iter(dataset))
# Normalize the image similar to DCN strategy
normalize = image / 255.
subtract_mean = tf.reduce_mean(normalize, axis=0)
x_train_from_our_dataloader = normalize - subtract_mean
x_train_from_dcn_dataloader = self.get_dcn_dataloader(self.image_white)
# compare output between our Dataloader and DCN dataloader
self.assertEqual(
np.allclose(x_train_from_our_dataloader,
x_train_from_dcn_dataloader), True)
# clean up
shutil.rmtree(dataset_dir)
def evaluate_dataset(args, model):
print(f"Evaluating on {args['dataloader']['name']}")
args['dataloader']['train_split_id'] = None
dataset = dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['val_list'],
num_classes=args["num_classes"],
split=args['dataloader']['val_split_id'])
model.evaluate(dataset)
def test_compare_with_augmentation_translate(self):
# create fake dataset
dataset_dir = create_fake_dataset_and_convert_to_tfrecords(
image_array=self.batch_random_images, n_images_per_class=10)
config = {
'name': 'tfrecords',
'data_dir': dataset_dir,
'batch_size': BATCH_SIZE,
'train_split_id': 'train',
'Translate': {
'width_shift_range': 0.125,
'height_shift_range': 0.125,
'padding_strategy': 'REFLECT'
}
}
dataset = get_dataset(config, ['Translate'], 1, 'train')
image, _ = next(iter(dataset))
# Normalize the image similar to DCN strategy
normalize = image / 255.
subtract_mean = tf.reduce_mean(normalize, axis=0)
x_train_from_our_dataloader = normalize - subtract_mean
# get x_train from dcn
x_data = self.get_dcn_dataloader(self.batch_random_images)
# DCN augmentation strategy
dcn_augmentations = keras_preprocessing.ImageDataGenerator(
height_shift_range=config['Translate']['height_shift_range'],
width_shift_range=config['Translate']['width_shift_range'],
horizontal_flip=False)
x_train_from_dcn_dataloader = next(
iter(
dcn_augmentations.flow(x_data,
None,
batch_size=BATCH_SIZE,
shuffle=False)))
# Test the tensor shape remains the same
self.assertTrue(x_train_from_dcn_dataloader.shape,
x_train_from_dcn_dataloader.shape)
# Note for translate we expect the mean to be different as the method to fill the missing pixels that keras and our data loader use are different.
# This doesn't cause an issue with the image being translated (see test_data_visualize), rather difference in pixel values at the shifted places.
# data visualization does confirm the translation works as expected.
self.assertAlmostEqual(np.mean(x_train_from_our_dataloader),
np.mean(x_train_from_dcn_dataloader),
places=3)
# clean up
shutil.rmtree(dataset_dir)
def train(args):
print(args)
global_conf.config_tf2(args)
checkpoint_dir, log_dir, export_dir = create_env_directories(
args, get_experiment_name(args))
train_dataset = dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['train_list'],
num_classes=args["num_classes"],
split=args['dataloader']['train_split_id'])
val_dataset = dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['val_list'],
num_classes=args["num_classes"],
split=args['dataloader']['val_split_id'])
setup_mp(args)
build_model_fn = get_model(args)
callbacks = get_callbacks(args, log_dir)
# tuner = Hyperband(build_model_fn,
# objective='val_accuracy',
# max_epochs=args['num_epochs'],
# hyperband_iterations=10e100,
# directory=checkpoint_dir)
tuner = BayesianOptimization(build_model_fn,
objective='val_accuracy',
max_trials=100000,
num_initial_points=10,
directory=checkpoint_dir)
tuner.search_space_summary()
tuner.search(x=train_dataset,
validation_data=val_dataset,
callbacks=callbacks,
epochs=args['num_epochs'])
tuner.results_summary()
def main():
args = argparse.parse_cmd(arguments)
datasets = {
'train':
dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['train_list'],
num_classes=10,
split='train'),
'val':
dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['val_list'],
num_classes=10,
split='test')
}
for dataset_type in ['train', 'val']:
for i, (images, y) in enumerate(datasets[dataset_type]):
image = images[0]
cv2.imwrite(os.path.join('/tmp', f'{dataset_type}_{i}.png'),
image.numpy())
if i == 20:
break
def train(args):
print(args)
global_conf.config_tf2(args)
checkpoint_dir, log_dir, export_dir = create_env_directories(
args, get_experiment_name(args))
train_dataset = dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['train_list'],
num_classes=args["num_classes"],
split=args['dataloader']['train_split_id'])
val_dataset = dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['val_list'],
num_classes=args["num_classes"],
split=args['dataloader']['val_split_id'])
setup_mp(args)
model, _ = define_model_in_strategy(args, get_model)
alchemy_api.send_model_info(model, args['server'])
callbacks = get_callbacks(args, log_dir)
model_checkpoint_cb, latest_epoch = init_custom_checkpoint_callbacks(
{'model': model}, checkpoint_dir)
callbacks.append(model_checkpoint_cb)
if args['server']['id'] != '':
callbacks.append(alchemy_api.send_metric_callbacks(args['server']))
model.fit(x=train_dataset,
validation_data=val_dataset,
epochs=args['num_epochs'],
callbacks=callbacks,
max_queue_size=16,
initial_epoch=latest_epoch)
print("export model")
export.export(model, export_dir, args)
print("Training Completed!!")
def train(args):
print(args)
global_conf.config_tf2(args)
checkpoint_dir, log_dir, export_dir = create_env_directories(args, get_experiment_name(args))
train_dataset = dataloader.get_dataset(args['dataloader'], transformation_list=args['dataloader']['train_list'],
num_classes=args["num_classes"], split=args['dataloader']['train_split_id'])
val_dataset = dataloader.get_dataset(args['dataloader'], transformation_list=args['dataloader']['val_list'],
num_classes=args["num_classes"], split=args['dataloader']['val_split_id'])
setup_mp(args)
model, _ = define_model_in_strategy(args, get_model)
alchemy_api.send_model_info(model, args['server'])
callbacks = get_callbacks(args, log_dir)
model_checkpoint_cb, latest_epoch = init_custom_checkpoint_callbacks({'model': model}, checkpoint_dir)
callbacks.append(model_checkpoint_cb)
if args['server']['id'] != '':
callbacks.append(alchemy_api.send_metric_callbacks(args['server']))
if args['model_name'] == 'Pdart':
from src.models.pdart import callback_epoch
callbacks.append(tf.keras.callbacks.LambdaCallback(on_epoch_begin=lambda epoch, logs: callback_epoch(epoch, args['num_epochs'], args['drop_path_prob'])))
# Use weight and biases only use_wandb is true and framework is tensorflow
if args['wandb_params']['use_wandb'] and 'tensorflow' in args['framework']:
from wandb.keras import WandbCallback
callbacks.append(WandbCallback())
model.fit(x=train_dataset,
validation_data=val_dataset,
epochs=args['num_epochs'],
callbacks=callbacks,
max_queue_size=16,
initial_epoch=latest_epoch
)
print("export model")
export.export(model, export_dir, args)
print("Training Completed!!")
def test_compare_without_augmentations_random(self):
# create fake dataset and convert TFrecords
dataset_dir = create_fake_dataset_and_convert_to_tfrecords(
image_array=self.batch_random_images, n_images_per_class=10)
# Code below tested with subtracting the mean for each color channels
# and there were differences between the output from our dataloader and DCN dataloader.
# For comparison purposes the mean is subtracted following DCN strategy.
# image_normalized = self.batch_random_images,/ 255.
# r = np.dstack([image_normalized[i][:, :, 0] for i in range(len(image_normalized))])
# g = np.dstack([image_normalized[i][:, :, 1] for i in range(len(image_normalized))])
# b = np.dstack([image_normalized[i][:, :, 2] for i in range(len(image_normalized))])
# mean_data = [np.mean(r), np.mean(g), np.mean(b)]
# print(mean_data)
config = {
'name': 'tfrecords',
'data_dir': dataset_dir,
'batch_size': BATCH_SIZE,
'train_split_id': 'train',
}
# get data from tfrecord
dataset = get_dataset(config, [], 1, 'train')
image, _ = next(iter(dataset))
# Normalize the image similar to DCN strategy
normalize = image / 255.
subtract_mean = tf.reduce_mean(normalize, axis=0)
x_train_from_our_dataloader = normalize - subtract_mean
# get x_train from dcn
x_train_from_dcn_dataloader = self.get_dcn_dataloader(
self.batch_random_images)
self.assertAlmostEqual(np.mean(x_train_from_our_dataloader),
np.mean(x_train_from_dcn_dataloader),
places=5)
# clean up
shutil.rmtree(dataset_dir)
def get_dataset(args):
args['dataloader']['train_split_id'] = None
dataset = dataloader.get_dataset(args['dataloader'], transformation_list=args['dataloader']['list'],
num_classes=args["num_classes"], split=args['dataloader']['split_id'])
return dataset
def train(config):
"""
This function setup:
1- Tensorflow (XLA, GPU configuration, mixed precision, execution strategies)
2- The datasets
3- The model
4- The execution environment
5- The monitoring (Upstride plateform and tensorboard)
Then
6- start the training
7- Export the model
"""
# 1
global_conf.config_tf2(config['config'])
global_conf.setup_mp(config['config'])
ds_strategy = global_conf.setup_strategy(config['config']['strategy'])
if config['model']['channels_first']: # if True set keras backend to channels_first
tf.keras.backend.set_image_data_format('channels_first')
# 2
train_dataset = dataloader.get_dataset(config['dataloader'], transformation_list=config['dataloader']['train_list'],
num_classes=config['model']["num_classes"], split=config['dataloader']['train_split_id'])
val_dataset = dataloader.get_dataset(config['dataloader'], transformation_list=config['dataloader']['val_list'],
num_classes=config['model']["num_classes"], split=config['dataloader']['val_split_id'])
# 3
with ds_strategy.scope():
model, optimizer = get_compiled_model(config)
# 4
checkpoint_dir, log_dir, export_dir = create_env_directories(get_experiment_name(config), config['checkpoint_dir'], config['log_dir'], config['export']['dir'])
if not os.path.exists(log_dir):
os.makedirs(log_dir)
with open(os.path.join(log_dir, "conf.yml"), 'w') as file:
yaml.dump(config, file)
# 5
config['server'] = alchemy_api.start_training(config['server'])
alchemy_api.send_model_info(model, config['server'])
callbacks = get_callbacks(config, log_dir)
with ds_strategy.scope(): # checkpoints needs to be in the same scope.
model_checkpoint_cb, latest_epoch = init_custom_checkpoint_callbacks({'model': model}, checkpoint_dir, config['max_checkpoints'], config['checkpoint_freq'])
callbacks.append(model_checkpoint_cb)
if config['server']['id'] != '':
callbacks.append(alchemy_api.send_metric_callbacks(config['server']))
if config['model']['name'] == 'Pdart':
from src.models.pdart import callback_epoch
callbacks.append(tf.keras.callbacks.LambdaCallback(on_epoch_begin=lambda epoch, logs: callback_epoch(epoch, config['num_epochs'], config['drop_path_prob'])))
# 6 training
if config['progressive_resizing']:
progressive_training(model=model,
config=config,
train_dataset=train_dataset,
val_dataset=val_dataset,
callbacks=callbacks,
latest_epoch=latest_epoch,
max_queue_size=16,
optimizer=optimizer)
else:
model.fit(x=train_dataset,
validation_data=val_dataset,
epochs=config['num_epochs'],
callbacks=callbacks,
max_queue_size=16,
initial_epoch=latest_epoch
)
# 7 training
print("export model")
export.export(model, export_dir, config)
print("Training Completed!!")
def test_data_visualization(self):
dataset_dir = tempfile.mkdtemp()
os.makedirs(os.path.join(dataset_dir, 'cat'), exist_ok=True)
shutil.copyfile(os.path.join(self.image_path),
os.path.join(dataset_dir, 'cat', 'cat.jpeg'))
image = keras_preprocessing.load_img(self.image_path,
grayscale=False,
color_mode='rgb',
target_size=None,
interpolation='nearest')
image_array = keras_preprocessing.img_to_array(image)
image_array = image_array[
np.newaxis, ...] # adding batch dimension for ImageDataLoader
args = {
'name': 'tfrecords',
'description': 'test',
'tfrecord_dir_path': dataset_dir,
'tfrecord_size': 1,
'preprocessing': 'NO',
'image_size': DIMENSIONS[1:3],
"n_tfrecords": 1,
'data': {
'images_dir_path': dataset_dir,
'annotation_file_path': None,
'delimiter': ',',
'header_exists': False,
'split_names': ['train'],
'split_percentages': [1.0],
}
}
# generate tfrecords
build_tfrecord_dataset(args)
config = {
'name': 'tfrecords',
'data_dir': dataset_dir,
'batch_size': 1,
'train_split_id': 'train',
'Translate': {
'width_shift_range': 0.125,
'height_shift_range': 0.125,
'padding_strategy': 'reflect'
}
}
# loop through until the RandomHorizontalFlip is applied.
dataset = get_dataset(config, ['Translate', 'RandomHorizontalFlip'], 1,
'train')
image, _ = next(iter(dataset)) # get the image
image_i = keras_preprocessing.array_to_img(
image[0]) # get the image excluding the batch dimension and save
keras_preprocessing.save_img(
os.path.join(dataset_dir, 'cat_after_augment_ours.png'),
image_i) # save the image
self.assertFalse(np.allclose(image, image_array),
"Augmented Image and original image are the same")
dcn_augmentations = keras_preprocessing.ImageDataGenerator(
height_shift_range=config['Translate']['height_shift_range'],
width_shift_range=config['Translate']['width_shift_range'],
horizontal_flip=True)
image_dcn = next(
iter(
dcn_augmentations.flow(image_array,
None,
batch_size=1,
shuffle=False)))
# save image excluding batch size
keras_preprocessing.save_img(
os.path.join(dataset_dir, 'cat_after_augment_DCN.png'), image[0])
self.assertFalse(np.allclose(image_dcn, image_array),
"Augmented Image and original image are the same")
# remove the below line in order to perform the visual inspection
shutil.rmtree(dataset_dir)
def test_compare_with_augmentation_randomflip(self):
# create fake dataset
# dataset_dir = create_fake_dataset_and_convert_to_tfrecords(image_array=self.black_white_image, n_images_per_class=1)
dataset_dir = tempfile.mkdtemp()
os.makedirs(os.path.join(dataset_dir, 'black_white'), exist_ok=True)
shutil.copyfile(
os.path.join(self.black_white_image),
os.path.join(dataset_dir, 'black_white', 'black_and_white.jpeg'))
image = keras_preprocessing.load_img(self.black_white_image,
grayscale=False,
color_mode='rgb',
target_size=None,
interpolation='nearest')
image_array = keras_preprocessing.img_to_array(image)
image_array = image_array[
np.newaxis, ...] # adding batch dimension for ImageDataLoader
args = {
'name': 'tfrecords',
'description': 'test',
'tfrecord_dir_path': dataset_dir,
'tfrecord_size': 1,
'preprocessing': 'NO',
'image_size': DIMENSIONS[1:3],
"n_tfrecords": 1,
'data': {
'images_dir_path': dataset_dir,
'annotation_file_path': None,
'delimiter': ',',
'header_exists': False,
'split_names': ['train'],
'split_percentages': [1.0],
}
}
# generate tfrecords
build_tfrecord_dataset(args)
config = {
'name': 'tfrecords',
'data_dir': dataset_dir,
'batch_size': 1,
'train_split_id': 'train',
}
# loop through until the RandomHorizontalFlip is applied
dataset = get_dataset(config, ['RandomHorizontalFlip'], 1, 'train')
x_train_from_our_dataloader, _ = next(iter(dataset))
self.assertFalse(np.allclose(x_train_from_our_dataloader, image_array),
"Augmented Image and original image are the same")
# loop through until the RandomHorizontalFlip is applied
dcn_augmentations = keras_preprocessing.ImageDataGenerator(
horizontal_flip=True)
x_train_from_dcn_dataloader = next(
iter(
dcn_augmentations.flow(image_array,
None,
batch_size=1,
shuffle=False)))
self.assertFalse(np.allclose(x_train_from_dcn_dataloader, image_array),
"Augmented Image and original image are the same")
# Test the tensor shape remains the same
self.assertTrue(x_train_from_our_dataloader.shape,
x_train_from_dcn_dataloader.shape)
# Note: The example used here is to test if the augmentation works using a black and white image.
# This test would fail if an image with wide range of pixel values across the color channels are used.
# This is a limitation due to the loss of data (encode to bytes and decode back to image) for JPEG format.
# In the TFrecord writer we use opencv and decode using tensorflow library. The encoded bytes using openCV and tensorflow are not identical.
# There is a slight pixel difference which is not noticible to the human eyes.
self.assertAlmostEqual(np.mean(x_train_from_our_dataloader),
np.mean(x_train_from_dcn_dataloader),
places=7)
# clean up
shutil.rmtree(dataset_dir)
def train(args):
# config_tf2(args['configuration']['xla'])
# Create log, checkpoint and export directories
checkpoint_dir, log_dir, export_dir = create_env_directories(
args, get_experiment_name(args))
train_log_dir = os.path.join(log_dir, 'train')
val_log_dir = os.path.join(log_dir, 'validation')
arch_log_dir = os.path.join(log_dir, 'arch')
summary_writers = {
'train': tf.summary.create_file_writer(train_log_dir),
'val': tf.summary.create_file_writer(val_log_dir),
'arch': tf.summary.create_file_writer(arch_log_dir)
}
# Prepare the 3 datasets
train_weight_dataset = dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['train_list'],
num_classes=args["num_classes"],
split='train_weights')
train_arch_dataset = dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['train_list'],
num_classes=args["num_classes"],
split='train_arch')
val_dataset = dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['val_list'],
num_classes=args["num_classes"],
split='test')
# define model, optimizer and checkpoint callback
setup_mp(args)
model = model_name_to_class[args['model_name']](
args['framework'],
input_shape=args['input_size'],
label_dim=args['num_classes']).model
model.summary()
alchemy_api.send_model_info(model, args['server'])
weights, arch_params = fbnetv2.split_trainable_weights(model)
weight_opt = get_optimizer(args['optimizer'])
arch_opt = get_optimizer(args['arch_search']['optimizer'])
model_checkpoint_cb, latest_epoch = init_custom_checkpoint_callbacks(
{'model': model}, checkpoint_dir, args['max_checkpoints'],
args['checkpoint_freq'])
callbacks = [model_checkpoint_cb]
temperature_decay_fn = fbnetv2.exponential_decay(
args['arch_search']['temperature']['init_value'],
args['arch_search']['temperature']['decay_steps'],
args['arch_search']['temperature']['decay_rate'])
lr_decay_fn = CosineDecay(
args['optimizer']['lr'],
alpha=args["optimizer"]["lr_decay_strategy"]["lr_params"]["alpha"],
total_epochs=args['num_epochs'])
lr_decay_fn_arch = CosineDecay(args['arch_search']['optimizer']['lr'],
alpha=0.000001,
total_epochs=args['num_epochs'])
metrics = {
'arch': {
'latency_reg_loss': tf.keras.metrics.Mean()
},
'train': {
'total_loss': tf.keras.metrics.Mean(),
'accuracy': tf.keras.metrics.CategoricalAccuracy(),
'cross_entropy_loss': tf.keras.metrics.Mean(),
},
'val': {
'accuracy': tf.keras.metrics.CategoricalAccuracy(),
'cross_entropy_loss': tf.keras.metrics.Mean(),
}
}
train_step = get_train_step_function(model, weights, weight_opt,
metrics['train'])
train_step_arch = get_train_step_arch_function(model, arch_params,
arch_opt, metrics['train'],
metrics['arch'])
evaluation_step = get_eval_step_function(model, metrics['val'])
for epoch in range(latest_epoch, args['num_epochs']):
print(f'Epoch: {epoch}/{args["num_epochs"]}')
# Update both LR
weight_opt.learning_rate = lr_decay_fn(epoch)
arch_opt.learning_rate = lr_decay_fn_arch(epoch)
# Updating the weight parameters using a subset of the training data
for step, (x_batch, y_batch) in tqdm.tqdm(
enumerate(train_weight_dataset, start=1)):
train_step(x_batch, y_batch)
# Evaluate the model on validation subset
for x_batch, y_batch in val_dataset:
evaluation_step(x_batch, y_batch)
# Handle metrics
template = f"Weights updated, Epoch {epoch}"
template = metrics_processing(metrics, summary_writers,
['train', 'val'], template, epoch)
template += f", lr: {float(weight_opt.learning_rate)}"
print(template)
new_temperature = temperature_decay_fn(epoch)
with summary_writers['train'].as_default():
tf.summary.scalar('temperature', new_temperature, step=epoch)
define_temperature(new_temperature)
if epoch >= args['arch_search']['num_warmup']:
# Updating the architectural parameters on another subset
for step, (x_batch, y_batch) in tqdm.tqdm(
enumerate(train_arch_dataset, start=1)):
train_step_arch(x_batch, y_batch)
# Evaluate the model on validation subset
for x_batch, y_batch in val_dataset:
evaluation_step(x_batch, y_batch)
# Handle metrics
template = f'Architecture updated, Epoch {epoch}'
template = metrics_processing(metrics,
summary_writers,
['train', 'val', 'arch'],
template,
epoch,
postfix='_arch')
template += f", lr: {float(arch_opt.learning_rate)}"
print(template)
# move saved outside of condition so we save starting from the begining
fbnetv2.save_arch_params(model, epoch, log_dir)
# manually call the callbacks
for callback in callbacks:
callback.on_epoch_end(epoch, logs=None)
print("Training Completed!!")
print("Architecture params: ")
print(arch_params)
fbnetv2.post_training_analysis(
model, args['arch_search']['exported_architecture'])
def train(args):
# config_tf2(args['configuration']['xla'])
# Create log, checkpoint and export directories
checkpoint_dir, log_dir, export_dir = create_env_directories(
args, get_experiment_name(args))
train_weight_dataset = dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['train_list'],
num_classes=args["num_classes"],
split='train_weights')
train_arch_dataset = dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['train_list'],
num_classes=args["num_classes"],
split='train_arch')
val_dataset = dataloader.get_dataset(
args['dataloader'],
transformation_list=args['dataloader']['val_list'],
num_classes=args["num_classes"],
split='validation')
setup_mp(args)
# define model, optimizer and checkpoint callback
model = model_name_to_class[args['model_name']](
args['framework'],
input_shape=args['input_size'],
label_dim=args['num_classes']).model
model.summary()
alchemy_api.send_model_info(model, args['server'])
weight_opt = get_optimizer(args['optimizer'])
arch_opt = get_optimizer(args['arch_optimizer_param'])
model_checkpoint_cb, latest_epoch = init_custom_checkpoint_callbacks(
{'model': model}, checkpoint_dir)
weights, arch_params = split_trainable_weights(model)
temperature_decay_fn = exponential_decay(
args['temperature']['init_value'], args['temperature']['decay_steps'],
args['temperature']['decay_rate'])
lr_decay_fn = CosineDecay(
args['optimizer']['lr'],
alpha=args["optimizer"]["lr_decay_strategy"]["lr_params"]["alpha"],
total_epochs=args['num_epochs'])
loss_fn = CategoricalCrossentropy()
accuracy_metric = CategoricalAccuracy()
loss_metric = Mean()
val_accuracy_metric = CategoricalAccuracy()
val_loss_metric = Mean()
train_log_dir = os.path.join(args['log_dir'], 'train')
val_log_dir = os.path.join(args['log_dir'], 'validation')
train_summary_writer = tf.summary.create_file_writer(train_log_dir)
val_summary_writer = tf.summary.create_file_writer(val_log_dir)
@tf.function
def train_step(x_batch, y_batch):
with tf.GradientTape() as tape:
y_hat = model(x_batch, training=True)
loss = loss_fn(y_batch, y_hat)
accuracy_metric.update_state(y_batch, y_hat)
loss_metric.update_state(loss)
grads = tape.gradient(loss, weights)
weight_opt.apply_gradients(zip(grads, weights))
@tf.function
def train_step_arch(x_batch, y_batch):
with tf.GradientTape() as tape:
y_hat = model(x_batch, training=False)
loss = loss_fn(y_batch, y_hat)
accuracy_metric.update_state(y_batch, y_hat)
loss_metric.update_state(loss)
grads = tape.gradient(loss, arch_params)
arch_opt.apply_gradients(zip(grads, arch_params))
@tf.function
def evaluation_step(x_batch, y_batch):
y_hat = model(x_batch, training=False)
loss = loss_fn(y_batch, y_hat)
val_accuracy_metric.update_state(y_batch, y_hat)
val_loss_metric.update_state(loss)
for epoch in range(latest_epoch, args['num_epochs']):
print(f'Epoch: {epoch}/{args["num_epochs"]}')
weight_opt.learning_rate = lr_decay_fn(epoch)
# Updating the weight parameters using a subset of the training data
for step, (x_batch, y_batch) in tqdm.tqdm(
enumerate(train_weight_dataset, start=1)):
train_step(x_batch, y_batch)
# Evaluate the model on validation subset
for x_batch, y_batch in val_dataset:
evaluation_step(x_batch, y_batch)
train_accuracy = accuracy_metric.result()
train_loss = loss_metric.result()
val_accuracy = val_accuracy_metric.result()
val_loss = val_loss_metric.result()
template = f'Weights updated, Epoch {epoch}, Train Loss: {float(train_loss)}, Train Accuracy: ' \
f'{float(train_accuracy)}, Val Loss: {float(val_loss)}, Val Accuracy: {float(val_accuracy)}, ' \
f'lr: {float(weight_opt.learning_rate)}'
print(template)
new_temperature = temperature_decay_fn(epoch)
with train_summary_writer.as_default():
tf.summary.scalar('loss', train_loss, step=epoch)
tf.summary.scalar('accuracy', train_accuracy, step=epoch)
tf.summary.scalar('temperature', new_temperature, step=epoch)
with val_summary_writer.as_default():
tf.summary.scalar('loss', val_loss, step=epoch)
tf.summary.scalar('accuracy', val_accuracy, step=epoch)
# Resetting metrices for reuse
accuracy_metric.reset_states()
loss_metric.reset_states()
val_accuracy_metric.reset_states()
val_loss_metric.reset_states()
if epoch >= 10:
# Updating the architectural parameters on another subset
for step, (x_batch, y_batch) in tqdm.tqdm(
enumerate(train_arch_dataset, start=1)):
train_step_arch(x_batch, y_batch)
# Evaluate the model on validation subset
for x_batch, y_batch in val_dataset:
evaluation_step(x_batch, y_batch)
train_accuracy = accuracy_metric.result()
train_loss = loss_metric.result()
val_accuracy = val_accuracy_metric.result()
val_loss = val_loss_metric.result()
template = f'Arch params updated, Epoch {epoch}, Train Loss: {float(train_loss)}, Train Accuracy: ' \
f'{float(train_accuracy)}, Val Loss: {float(val_loss)}, Val Accuracy: {float(val_accuracy)}'
print(template)
with train_summary_writer.as_default():
tf.summary.scalar('loss_after_arch_params_update',
train_loss,
step=epoch)
tf.summary.scalar('accuracy_after_arch_params_update',
train_accuracy,
step=epoch)
with val_summary_writer.as_default():
tf.summary.scalar('loss_after_arch_params_update',
val_loss,
step=epoch)
tf.summary.scalar('accuracy_after_arch_params_update',
val_accuracy,
step=epoch)
# Resetting metrices for reuse
accuracy_metric.reset_states()
loss_metric.reset_states()
val_accuracy_metric.reset_states()
val_loss_metric.reset_states()
define_temperature(new_temperature)
print("Training Completed!!")
print("Architecture params: ")
print(arch_params)
post_training_analysis(model, args['exported_architecture'])