def get_compiled_model(config): # Parameters usefull for all the models kwargs = { 'input_size': config['model']['input_size'], 'changing_ids': config['model']['changing_ids'], 'num_classes': config['model']['num_classes'], 'factor': config['model']['factor'], 'upstride_type': config['model']['upstride_type'], 'tf2upstride_strategy': config['model']['conversion_params']['tf2up_strategy'], 'upstride2tf_strategy': config['model']['conversion_params']['up2tf_strategy'], 'weight_decay': config['optimizer']['weight_decay'], } # for architecture search if config['load_searched_arch']: kwargs['load_searched_arch'] = config['load_searched_arch'] model = model_name_to_class[config['model']['name']](**kwargs).build() model.summary() # calculates FLOPs for the Model. if config['model']['calculate_flops']: calc_flops = metrics.count_flops_efficient(model, config['model']['upstride_type']) print(f"Total FLOPs for {config['model']['name']}: {calc_flops}") optimizer = get_optimizer(config['optimizer']) model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy', 'top_k_categorical_accuracy']) # output the optimizer to save it in the checkpoint return model, optimizer
def get_model(args): model = model_name_to_class[args['model_name']]( args['framework'], args['factor'], args['input_size'], args['num_classes'], args['n_layers_before_tf'], False).model model.summary() optimizer = get_optimizer(args['optimizer']) model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy', 'top_k_categorical_accuracy']) # output the optimizer to save it in the checkpoint return model, optimizer
def build_model(hp): factor = hp.Choice('factor', get_values_from_args(args['factor']), ordered=True) framework = hp.Choice('framework', args['frameworks']) model = model_name_to_class[args['model_name']](framework, factor, args['input_size'], args['num_classes'], hp=hp).model model.compile(optimizer=get_optimizer(args['optimizer']), loss='categorical_crossentropy', metrics=['accuracy']) return model
def get_model(args): load_arch = args['load_searched_arch'] if args['load_searched_arch'] else None model = model_name_to_class[args['model_name']](args['framework'], # TODO replace args[] by args['model'] args['conversion_params'], args['factor'], args['input_size'], args['num_classes'], args['n_layers_before_tf'], False, load_searched_arch=load_arch, args=args).model model.summary() optimizer = get_optimizer(args['optimizer']) model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy', 'top_k_categorical_accuracy']) # output the optimizer to save it in the checkpoint return model, optimizer
def progressive_training(model, config, train_dataset, val_dataset, callbacks, latest_epoch, max_queue_size=16, optimizer=None): """Trains the model using progressive resizing Images resolution will be modified during the training following config["resizing_sizes"] The specific portions of the total number of epochs dedicated to each size is defined in config["resizing_training_portions"] example: image resolution = 768 x 512 config["resizing_sizes"] = [0.5, 1] config["resizing_training_portions"] = [0.7,0.3] config['num_epochs'] = 100 -> For the 0.7*100=70 first epochs the network will train with (0.5*768 x 0.5*512) = (384 x 256) resolution For the 0.3*100=30 following epochs, the network will train with (1*768 x 1*512) = (768 x 512) resolution if config["reset_after_resize"] is True, the callbacks and optimizer will be reset after each change in resolution. """ sizes_list = config["resizing_sizes"] resizing_portions = config["resizing_training_portions"] assert len(sizes_list) == len(resizing_portions), "resizing_sizes and resizing_portions must have the same length" assert sum(resizing_portions) == 1, "the sum of resizing_portions must equal 1" reset = config["reset_after_resize"] if reset: callbacks_tmp = [copy.deepcopy(callbacks) for _ in range(len(sizes_list))] total_epochs = config['num_epochs'] epochs_by_portion = [int(total_epochs * resizing_portion) for resizing_portion in resizing_portions] epochs_by_portion[-1] = total_epochs - sum(epochs_by_portion[:-1]) # to ensure the right total nb of epochs current_epoch = 0 for element in train_dataset: original_size = element[0].shape # N, H, W, C break for i, size in enumerate(sizes_list): if latest_epoch >= current_epoch + epochs_by_portion[i]: current_epoch += epochs_by_portion[i] continue current_epoch = max(latest_epoch, current_epoch) if abs(size-1) > 0.001: resize = tf.keras.layers.experimental.preprocessing.Resizing(int(original_size[1] * size), int(original_size[2] * size)) resized_train_set = train_dataset.map(lambda x, y: (resize(x), y)).prefetch(tf.data.experimental.AUTOTUNE) else: resized_train_set = train_dataset if reset: callbacks = callbacks_tmp[i] optimizer = get_optimizer(config['optimizer']) model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy', 'top_k_categorical_accuracy']) model.fit(x=resized_train_set, validation_data=val_dataset, epochs=current_epoch + epochs_by_portion[i], callbacks=callbacks, max_queue_size=max_queue_size, initial_epoch=current_epoch ) current_epoch += epochs_by_portion[i]
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'])