loss = lossobj(masks, preds) valid_loss(loss) valid_accuracy(masks, preds) tloss = [] vloss = [] taccuracy = [] vaccuracy = [] bestloss = sys.maxsize # training through epochs for epoch in range(args.epochs): train_loss.reset_states() train_accuracy.reset_states() if not args.train: valid_loss.reset_states() valid_accuracy.reset_states() for imgs, masks in train_data: trainStep(imgs, masks) for imgs, masks in valid_data: if not args.train: validStep(imgs, masks) else: trainStep(imgs, masks) template1 = "epoch[{0}/{1}] training: mean loss: {2}, accuracy: {3}"
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'])