File size: 8,579 Bytes
a6b4c35 af2f07c a6b4c35 af2f07c a6b4c35 af2f07c a6b4c35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | '''
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
"""
# set data index
offset_output_index=0
input_index=1
output_index=2
# initialization
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): # opt.niter + opt.niter_decay + 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
# print(f"dec_outputs size is : {dec_outputs.shape}")
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()
# set data index
offset_output_index=0
input_index=1
output_index=2
# create containers
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()
# get model outputs
gen_outputs = model.decoder.output # gen_outputs.size = (N,1,H,W)
if opt_exp.n_decoders == 2:
off_outputs = model.offset_decoder.output # off_outputs.size = (N,n_ap,H,W)
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 # default save directory
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
|