| '''
|
| Scripts for the training and testing functions
|
| train() function is called for training the network
|
| test() function is called to evaluate the network
|
| Both the function logs and saves the results in the files
|
| as mentioned in the params.py file
|
| '''
|
| import warnings
|
| with warnings.catch_warnings():
|
| warnings.filterwarnings("ignore",category=FutureWarning)
|
| import numpy as np
|
| import os
|
| import time
|
| import hdf5storage
|
| from utils import *
|
| from Generators import *
|
| from params import *
|
|
|
| def train(model, train_loader, test_loader):
|
| """Traning pipeline
|
|
|
| Args:
|
| model (torch.module): pytorch model
|
| train_loader (torch.dataloader): dataloader
|
| test_loader (torch.dataloader): dataloader
|
| """
|
|
|
| offset_output_index=0
|
| input_index=1
|
| output_index=2
|
|
|
|
|
| total_steps = 0
|
| print('Training called')
|
| best_median = float('inf')
|
|
|
| for epoch in range(model.opt.starting_epoch_count+1, model.opt.n_epochs+1):
|
| epoch_start_time = time.time()
|
| epoch_loss = 0
|
| epoch_offset_loss = 0
|
| error =[]
|
|
|
| for i, data in enumerate(train_loader):
|
| total_steps += model.opt.batch_size
|
| if opt_exp.n_decoders == 2:
|
| model.set_input(data[input_index], data[output_index], data[offset_output_index], shuffle_channel=False)
|
| elif opt_exp.n_decoders == 1:
|
| model.set_input(data[input_index], data[output_index], shuffle_channel=False)
|
| model.optimize_parameters()
|
| dec_outputs = model.decoder.output
|
|
|
| error.extend(localization_error(dec_outputs.data.cpu().numpy(),data[output_index].cpu().numpy(),scale=0.1))
|
|
|
| write_log([str(model.decoder.loss.item())], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='loss')
|
| if opt_exp.n_decoders == 2:
|
| write_log([str(model.offset_decoder.loss.item())], model.offset_decoder.model_name, log_dir=model.offset_decoder.opt.log_dir, log_type='offset_loss')
|
| if total_steps % model.decoder.opt.save_latest_freq == 0:
|
| print('saving the latest model (epoch %d, total_steps %d)' %
|
| (epoch, total_steps))
|
| model.save_networks('latest')
|
|
|
| epoch_loss += model.decoder.loss.item()
|
| if opt_exp.n_decoders == 2:
|
| epoch_offset_loss += model.offset_decoder.loss.item()
|
|
|
| median_error_tr = np.median(error)
|
| error_90th_tr = np.percentile(error,90)
|
| error_99th_tr = np.percentile(error,99)
|
| nighty_percentile_error_tr = np.percentile(error,90)
|
| epoch_loss /= i
|
| if opt_exp.n_decoders == 2:
|
| epoch_offset_loss /= i
|
| write_log([str(epoch_loss)], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='epoch_decoder_loss')
|
| if opt_exp.n_decoders == 2:
|
| write_log([str(epoch_offset_loss)], model.offset_decoder.model_name, log_dir=model.offset_decoder.opt.log_dir, log_type='epoch_offset_decoder_loss')
|
| write_log([str(median_error_tr)], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='train_median_error')
|
| write_log([str(error_90th_tr)], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='train_90th_error')
|
| write_log([str(error_99th_tr)], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='train_99th_error')
|
| write_log([str(nighty_percentile_error_tr)], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='train_90_error')
|
| if (epoch == model.opt.starting_epoch_count + 1):
|
| min_eval_loss, median_error = test(model, test_loader, save_output=False)
|
| best_median = median_error
|
| print('Saving initial best model at epoch %d (median=%.4f)' % (epoch, best_median))
|
| model.save_networks('best')
|
| else:
|
| new_eval_loss, new_med_error = test(model, test_loader, save_output=False)
|
| if new_med_error < best_median:
|
| best_median = new_med_error
|
| print('Saving new best model at epoch %d (median=%.4f)' % (epoch, best_median))
|
| model.save_networks('best')
|
| median_error = new_med_error
|
|
|
| if epoch % model.encoder.opt.save_epoch_freq == 0:
|
| print('saving the model at the end of epoch %d, iters %d' %(epoch, total_steps))
|
| model.save_networks('latest')
|
| model.save_networks(epoch)
|
|
|
| print('End of epoch %d / %d \t Time Taken: %d sec' %
|
| (epoch, model.decoder.opt.niter + model.decoder.opt.niter_decay, time.time() - epoch_start_time))
|
| model.decoder.update_learning_rate()
|
| model.encoder.update_learning_rate()
|
| if opt_exp.n_decoders == 2:
|
| model.offset_decoder.update_learning_rate()
|
|
|
|
|
| def test(model, test_loader, save_output=True, save_name="decoder_test_result", save_dir="", log=True):
|
| """Test and evaluation pipeline
|
|
|
| Args:
|
| model (torch.module): pytorch model
|
| test_loader (torch.dataloader): dataloader
|
| save_output (bool, optional): whether to save output to mat file. Defaults to True.
|
| save_name (str, optional): name of the mat file. Defaults to "decoder_test_result".
|
| save_dir (str, optional): directory where output mat file is saved. Defaults to "".
|
| log (bool, optional): whether to log output. Defaults to True.
|
|
|
| Returns:
|
| tuple: (total_loss -> float, median_error -> float)
|
| """
|
| print('Evaluation Called')
|
| model.eval()
|
|
|
|
|
| offset_output_index=0
|
| input_index=1
|
| output_index=2
|
|
|
|
|
| generated_outputs = []
|
| offset_outputs = []
|
| total_loss = 0
|
| total_offset_loss = 0
|
| error =[]
|
| for i, data in enumerate(test_loader):
|
| if opt_exp.n_decoders == 2:
|
| model.set_input(data[input_index], data[output_index], data[offset_output_index], shuffle_channel=False)
|
| elif opt_exp.n_decoders == 1:
|
| model.set_input(data[input_index], data[output_index], shuffle_channel=False)
|
| model.test()
|
|
|
|
|
| gen_outputs = model.decoder.output
|
| if opt_exp.n_decoders == 2:
|
| off_outputs = model.offset_decoder.output
|
|
|
| generated_outputs.extend(gen_outputs.data.cpu().numpy())
|
| if opt_exp.n_decoders == 2:
|
| offset_outputs.extend(off_outputs.data.cpu().numpy())
|
| error.extend(localization_error(gen_outputs.data.cpu().numpy(),data[output_index].cpu().numpy(),scale=0.1))
|
| total_loss += model.decoder.loss.item()
|
| if opt_exp.n_decoders == 2:
|
| total_offset_loss += model.offset_decoder.loss.item()
|
| total_loss /= i
|
| if opt_exp.n_decoders == 2:
|
| total_offset_loss /= i
|
| median_error = np.median(error)
|
| nighty_percentile_error = np.percentile(error,90)
|
| error_99th = np.percentile(error,99)
|
|
|
| if log:
|
| write_log([str(median_error)], model.decoder.model_name, log_dir=model.opt.log_dir, log_type='test_median_error')
|
| write_log([str(nighty_percentile_error)], model.decoder.model_name, log_dir=model.opt.log_dir, log_type='test_90_error')
|
| write_log([str(error_99th)], model.decoder.model_name, log_dir=model.opt.log_dir, log_type='test_99_error')
|
| write_log([str(total_loss)], model.decoder.model_name, log_dir=model.opt.log_dir, log_type='test_loss')
|
| if opt_exp.n_decoders == 2:
|
| write_log([str(total_offset_loss)], model.offset_decoder.model_name, log_dir=model.opt.log_dir, log_type='test_offset_loss')
|
|
|
| if save_output:
|
| if not save_dir:
|
| save_dir = model.decoder.results_save_dir
|
|
|
| if not os.path.exists(save_dir):
|
| os.makedirs(save_dir, exist_ok=True)
|
|
|
| save_path = f"{save_dir}/{save_name}.mat"
|
| hdf5storage.savemat(save_path,
|
| mdict={"outputs":generated_outputs,"wo_outputs":offset_outputs, "error": error},
|
| appendmat=True,
|
| format='7.3',
|
| truncate_existing=True)
|
| print(f"result saved in {save_path}")
|
| return total_loss, median_error
|
|
|
