def make_model(src_vocab, tgt_vocab, emb_size=256, hidden_size=512, num_layers=1, dropout=0.1): "Helper: Construct a model from hyperparameters." attention = model.BahdanauAttention(hidden_size) mdl = model.EncoderDecoder( model.Encoder(emb_size, hidden_size, num_layers=num_layers, dropout=dropout), model.Decoder(emb_size, hidden_size, attention, num_layers=num_layers, dropout=dropout), nn.Embedding(src_vocab, emb_size), nn.Embedding(tgt_vocab, emb_size), model.Generator(hidden_size, tgt_vocab)) return mdl.cuda() if USE_CUDA else mdl
def loadModel(path): encoderDecoder = model.EncoderDecoder() encoderDecoder.load_state_dict(torch.load(path)) return encoderDecoder
def random_cv(cv_index, cv_year, roothpath, param_grid, num_random, model_name, device, one_day): """Hyperparameter tuning through random search Args: cv_index: the month of the valiation set cv_year: the year of the valiation set rootpath: the path where training-validtion sets are saved param_grid: a dictionary, consisting the grid of hyperparameters num_randon: the number of sets of hyperparameters to evaluate(tune) model_name: a string representing the name of a model device: indicates if the model should be run on cpu or gpu one_day: True or False, indicating if only the most recent available day is used for training a model (XGBoost or Lasso) """ # load data if model_name in ['CNN_LSTM', 'CNN_FNN']: train_X = joblib.load( rootpath + 'train_X_map_{}_forecast{}.pkl'.format(cv_year, cv_index)) valid_X = joblib.load( rootpath + 'val_X_map_{}_forecast{}.pkl'.format(cv_year, cv_index)) train_y = load_results( rootpath + 'train_y_pca_{}_forecast{}.pkl'.format(cv_year, cv_index)) valid_y = load_results( rootpath + 'val_y_pca_{}_forecast{}.pkl'.format(cv_year, cv_index)) output_dim = train_y.shape[-1] else: train_X = load_results( rootpath + 'train_X_pca_{}_forecast{}.pkl'.format(cv_year, cv_index)) valid_X = load_results( rootpath + 'val_X_pca_{}_forecast{}.pkl'.format(cv_year, cv_index)) train_y = load_results( rootpath + 'train_y_pca_{}_forecast{}.pkl'.format(cv_year, cv_index)) valid_y = load_results( rootpath + 'val_y_pca_{}_forecast{}.pkl'.format(cv_year, cv_index)) # set input and output dim input_dim = train_X.shape[-1] output_dim = train_y.shape[-1] if model_name == 'EncoderFNN_AllSeq_AR_CI' or model_name == 'EncoderFNN_AllSeq_AR': hidden_dim = param_grid['hidden_dim'] num_layers = param_grid['num_layers'] lr = param_grid['learning_rate'] threshold = param_grid['threshold'] num_epochs = param_grid['num_epochs'] seq_len = param_grid['seq_len'] linear_dim = param_grid['linear_dim'] drop_out = param_grid['drop_out'] if model_name == 'EncoderFNN_AllSeq_AR_CI': ci_dim = param_grid['ci_dim'] train_y_ar = load_results( rootpath + 'train_y_pca_ar_{}_forecast{}.pkl'.format(cv_year, cv_index)) valid_y_ar = load_results( rootpath + 'val_y_pca_ar_{}_forecast{}.pkl'.format(cv_year, cv_index)) train_dataset = model.MapDataset_ar(train_X, train_y_ar, train_y) train_loader = DataLoader(dataset=train_dataset, batch_size=512, shuffle=False) elif model_name == 'EncoderDecoder' or model_name == 'EncoderFNN_AllSeq' or model_name == 'EncoderFNN': train_dataset = model.MapDataset(train_X, train_y) train_loader = DataLoader(dataset=train_dataset, batch_size=512, shuffle=False) hidden_dim = param_grid['hidden_dim'] num_layers = param_grid['num_layers'] lr = param_grid['learning_rate'] threshold = param_grid['threshold'] num_epochs = param_grid['num_epochs'] if model_name == 'EncoderDecoder': decoder_len = param_grid['decoder_len'] elif model_name == 'EncoderFNN': last_layer = param_grid['last_layer'] seq_len = param_grid['seq_len'] elif model_name == 'EncoderFNN_AllSeq': seq_len = param_grid['seq_len'] linear_dim = param_grid['linear_dim'] drop_out = param_grid['drop_out'] elif model_name == 'XGBoost': if one_day is True: train_X = train_X[:, -1, :] # one day valid_X = valid_X[:, -1, :] # one day train_X = np.reshape(train_X, (train_X.shape[0], -1)) valid_X = np.reshape(valid_X, (valid_X.shape[0], -1)) max_depth = param_grid['max_depth'] colsample_bytree = param_grid['colsample_bytree'] gamma = param_grid['gamma'] n_estimators = param_grid['n_estimators'] lr = param_grid['learning_rate'] elif model_name == 'Lasso': if one_day is True: train_X = train_X[:, -1, :] # one day valid_X = valid_X[:, -1, :] # one day train_X = np.reshape(train_X, (train_X.shape[0], -1)) valid_X = np.reshape(valid_X, (valid_X.shape[0], -1)) alphas = param_grid['alpha'] elif model_name == 'FNN': if one_day is True: train_X = train_X[:, -1, :] # one day valid_X = valid_X[:, -1, :] # one day train_X = np.reshape(train_X, (train_X.shape[0], -1)) valid_X = np.reshape(valid_X, (valid_X.shape[0], -1)) train_dataset = model.MapDataset(train_X, train_y) train_loader = DataLoader(dataset=train_dataset, batch_size=512, shuffle=False) hidden_dim = param_grid['hidden_dim'] num_layers = param_grid['num_layers'] elif model_name == 'CNN_FNN': train_dataset = model.MapDataset_CNN(train_X, train_y) train_loader = DataLoader(dataset=train_dataset, batch_size=50, shuffle=False) stride = param_grid['stride'] kernel_size = param_grid['kernel_size'] hidden_dim = param_grid['hidden_dim'] num_layers = param_grid['num_layers'] elif model_name == 'CNN_LSTM': train_dataset = model.MapDataset_CNN(train_X, train_y) train_loader = DataLoader(dataset=train_dataset, batch_size=50, shuffle=False) stride = param_grid['module__stride'] kernel_size = param_grid['module__kernel_size'] hidden_dim = param_grid['module__hidden_dim'] num_lstm_layers = param_grid['module__num_lstm_layers'] lr = param_grid['lr'] num_epochs = param_grid['module__num_epochs'] else: print('the model name is not in the list') history_all = [] score = [] parameter_all = [] for i in range(num_random): # set model if model_name == 'EncoderDecoder': curr_hidden_dim = hidden_dim[randint(0, len(hidden_dim) - 1)] curr_num_layer = num_layers[randint(0, len(num_layers) - 1)] curr_decoder_len = decoder_len[randint(0, len(decoder_len) - 1)] curr_threshold = threshold[randint(0, len(threshold) - 1)] curr_lr = lr[randint(0, len(lr) - 1)] curr_num_epochs = num_epochs[randint(0, len(num_epochs) - 1)] parameters = { 'hidden_dim': curr_hidden_dim, 'num_layers': curr_num_layer, 'decoder_len': curr_decoder_len, 'threshold': curr_threshold, 'learning_rate': curr_lr, 'num_epochs': curr_num_epochs } parameter_all.append(parameters) mdl = model.EncoderDecoder(input_dim=input_dim, output_dim=output_dim, hidden_dim=curr_hidden_dim, num_layers=curr_num_layer, learning_rate=curr_lr, decoder_len=curr_decoder_len, threshold=curr_threshold, num_epochs=curr_num_epochs) # initialize the model model.init_weight(mdl) # send model to gpu mdl.to(device) # fit the model history = mdl.fit_cv(train_loader, valid_X, valid_y, device) # compute the prediction of validation set pred_y = mdl.predict(valid_X, device) elif model_name == 'EncoderFNN': curr_hidden_dim = hidden_dim[randint(0, len(hidden_dim) - 1)] curr_num_layer = num_layers[randint(0, len(num_layers) - 1)] curr_seq_len = seq_len[randint(0, len(seq_len) - 1)] curr_threshold = threshold[randint(0, len(threshold) - 1)] curr_lr = lr[randint(0, len(lr) - 1)] curr_num_epochs = num_epochs[randint(0, len(num_epochs) - 1)] curr_last_layer = last_layer[randint(0, len(last_layer) - 1)] parameters = { 'hidden_dim': curr_hidden_dim, 'num_layers': curr_num_layer, 'last_layer': curr_last_layer, 'threshold': curr_threshold, 'learning_rate': curr_lr, 'num_epochs': curr_num_epochs, 'seq_len': curr_seq_len } parameter_all.append(parameters) mdl = model.EncoderFNN(input_dim=input_dim, output_dim=output_dim, hidden_dim=curr_hidden_dim, num_layers=curr_num_layer, last_layer=curr_last_layer, seq_len=curr_seq_len, learning_rate=curr_lr, threshold=curr_threshold, num_epochs=curr_num_epochs) # initialize the model model.init_weight(mdl) # send model to gpu mdl.to(device) # fit the model history = mdl.fit_cv(train_loader, valid_X, valid_y, device) # compute the prediction of validation set pred_y = mdl.predict(valid_X, device) elif model_name == 'EncoderFNN_AllSeq': curr_hidden_dim = hidden_dim[randint(0, len(hidden_dim) - 1)] curr_num_layer = num_layers[randint(0, len(num_layers) - 1)] curr_seq_len = seq_len[randint(0, len(seq_len) - 1)] curr_threshold = threshold[randint(0, len(threshold) - 1)] curr_lr = lr[randint(0, len(lr) - 1)] curr_num_epochs = num_epochs[randint(0, len(num_epochs) - 1)] curr_linear_dim = linear_dim[randint(0, len(linear_dim) - 1)] curr_drop_out = drop_out[randint(0, len(drop_out) - 1)] parameters = { 'hidden_dim': curr_hidden_dim, 'num_layers': curr_num_layer, 'linear_dim': curr_linear_dim, 'threshold': curr_threshold, 'learning_rate': curr_lr, 'num_epochs': curr_num_epochs, 'seq_len': curr_seq_len, 'drop_out': curr_drop_out } parameter_all.append(parameters) mdl = model.EncoderFNN_AllSeq(input_dim=input_dim, output_dim=output_dim, hidden_dim=curr_hidden_dim, num_layers=curr_num_layer, seq_len=curr_seq_len, linear_dim=curr_linear_dim, learning_rate=curr_lr, dropout=curr_drop_out, threshold=curr_threshold, num_epochs=curr_num_epochs) # initialize the model model.init_weight(mdl) # send model to gpu mdl.to(device) # fit the model history = mdl.fit_cv(train_loader, valid_X, valid_y, device) # compute the prediction of validation set pred_y = mdl.predict(valid_X, device) elif model_name == 'EncoderFNN_AllSeq_AR': curr_hidden_dim = hidden_dim[randint(0, len(hidden_dim) - 1)] curr_num_layer = num_layers[randint(0, len(num_layers) - 1)] curr_seq_len = seq_len[randint(0, len(seq_len) - 1)] curr_threshold = threshold[randint(0, len(threshold) - 1)] curr_lr = lr[randint(0, len(lr) - 1)] curr_num_epochs = num_epochs[randint(0, len(num_epochs) - 1)] curr_linear_dim = linear_dim[randint(0, len(linear_dim) - 1)] curr_drop_out = drop_out[randint(0, len(drop_out) - 1)] parameters = { 'hidden_dim': curr_hidden_dim, 'num_layers': curr_num_layer, 'linear_dim': curr_linear_dim, 'threshold': curr_threshold, 'learning_rate': curr_lr, 'num_epochs': curr_num_epochs, 'seq_len': curr_seq_len, 'drop_out': curr_drop_out } parameter_all.append(parameters) mdl = model.EncoderFNN_AllSeq_AR(input_dim=input_dim, output_dim=output_dim, hidden_dim=curr_hidden_dim, num_layers=curr_num_layer, seq_len=curr_seq_len, linear_dim=curr_linear_dim, learning_rate=curr_lr, dropout=curr_drop_out, threshold=curr_threshold, num_epochs=curr_num_epochs) # initialize the model model.init_weight(mdl) # send model to gpu mdl.to(device) # fit the model history = mdl.fit_cv(train_loader, valid_X, valid_y_ar, valid_y, device) # compute the prediction of validation set pred_y = mdl.predict(valid_X, valid_y_ar, device) elif model_name == 'EncoderFNN_AllSeq_AR_CI': curr_hidden_dim = hidden_dim[randint(0, len(hidden_dim) - 1)] curr_num_layer = num_layers[randint(0, len(num_layers) - 1)] curr_seq_len = seq_len[randint(0, len(seq_len) - 1)] curr_threshold = threshold[randint(0, len(threshold) - 1)] curr_lr = lr[randint(0, len(lr) - 1)] curr_num_epochs = num_epochs[randint(0, len(num_epochs) - 1)] curr_linear_dim = linear_dim[randint(0, len(linear_dim) - 1)] curr_drop_out = drop_out[randint(0, len(drop_out) - 1)] parameters = { 'hidden_dim': curr_hidden_dim, 'num_layers': curr_num_layer, 'linear_dim': curr_linear_dim, 'threshold': curr_threshold, 'learning_rate': curr_lr, 'num_epochs': curr_num_epochs, 'seq_len': curr_seq_len, 'drop_out': curr_drop_out, 'ci_dim': ci_dim } parameter_all.append(parameters) mdl = model.EncoderFNN_AllSeq_AR_CI(input_dim=input_dim - ci_dim, output_dim=output_dim, hidden_dim=curr_hidden_dim, num_layers=curr_num_layer, seq_len=curr_seq_len, linear_dim=curr_linear_dim, ci_dim=ci_dim, learning_rate=curr_lr, dropout=curr_drop_out, threshold=curr_threshold, num_epochs=curr_num_epochs) # initialize the model model.init_weight(mdl) # send model to gpu mdl.to(device) # fit the model history = mdl.fit_cv(train_loader, valid_X, valid_y_ar, valid_y, device) pred_y = mdl.predict(valid_X, valid_y_ar, device) elif model_name == 'XGBoost': curr_max_depth = max_depth[randint(0, len(max_depth) - 1)] curr_colsample_bytree = colsample_bytree[randint( 0, len(colsample_bytree) - 1)] curr_gamma = gamma[randint(0, len(gamma) - 1)] curr_n_estimators = n_estimators[randint(0, len(n_estimators) - 1)] curr_lr = lr[randint(0, len(lr) - 1)] parameters = { 'max_depth': curr_max_depth, 'colsample_bytree': curr_colsample_bytree, 'gamma': curr_gamma, 'n_estimators': curr_n_estimators, 'learning_rate': curr_lr } parameter_all.append(parameters) mdl = model.XGBMultitask(num_models=output_dim, colsample_bytree=curr_colsample_bytree, gamma=curr_gamma, learning_rate=curr_lr, max_depth=curr_max_depth, n_estimators=curr_n_estimators, objective='reg:squarederror') # history = mdl.fit_cv(train_X, train_y, valid_X, valid_y) mdl.fit(train_X, train_y) pred_y = mdl.predict(valid_X) history = None elif model_name == 'Lasso': curr_alpha = alphas[randint(0, len(alphas) - 1)] parameter = {'alpha': curr_alpha} parameter_all.append(parameter) mdl = model.LassoMultitask(alpha=curr_alpha, fit_intercept=False) mdl.fit(train_X, train_y) pred_y = mdl.predict(valid_X) history = None elif model_name == 'FNN': curr_hidden_dim = hidden_dim[randint(0, len(hidden_dim) - 1)] curr_num_layers = num_layers[randint(0, len(num_layers) - 1)] parameters = { 'hidden_dim': curr_hidden_dim, 'num_layers': curr_num_layers } parameter_all.append(parameters) mdl = model.ReluNet(input_dim=input_dim, output_dim=output_dim, hidden_dim=curr_hidden_dim, num_layers=curr_num_layers, threshold=0.1, num_epochs=1000) model.init_weight(mdl) mdl.to(device) history = mdl.fit_cv(train_loader, valid_X, valid_y, device) pred_y = mdl.predict(valid_X, device) elif model_name == 'CNN_FNN': curr_stride = stride[randint(0, len(stride) - 1)] curr_kernel_size = kernel_size[randint(0, len(kernel_size) - 1)] curr_hidden_dim = hidden_dim[randint(0, len(hidden_dim) - 1)] curr_num_layers = num_layers[randint(0, len(num_layers) - 1)] parameters = { 'stride': curr_stride, 'kernel_size': curr_kernel_size, 'hidden_dim': curr_hidden_dim, 'num_layers': curr_num_layers } parameter_all.append(parameters) num_var = len(train_X) input_dim = model.get_input_dim(train_X, num_var, curr_stride, curr_kernel_size) mdl = model.CnnFnn(num_var, input_dim, output_dim, kernel_size=curr_kernel_size, stride=curr_stride, hidden_dim=curr_hidden_dim, num_layers=curr_num_layers, num_epochs=100) mdl.to(device) history = mdl.fit_cv(train_loader, valid_X, valid_y, device) pred_y = mdl.predict(valid_X, device) elif model_name == 'CNN_LSTM': curr_stride = stride[randint(0, len(stride) - 1)] curr_kernel_size = kernel_size[randint(0, len(kernel_size) - 1)] curr_hidden_dim = hidden_dim[randint(0, len(hidden_dim) - 1)] curr_num_layers = num_lstm_layers[randint(0, len(num_lstm_layers) - 1)] curr_lr = lr[randint(0, len(lr) - 1)] curr_num_epochs = num_epochs[randint(0, len(num_epochs) - 1)] parameters = { 'stride': curr_stride, 'kernel_size': curr_kernel_size, 'hidden_dim': curr_hidden_dim, 'num_layers': curr_num_layers, 'learning_rate': curr_lr, 'num_epochs': curr_num_epochs } parameter_all.append(parameters) num_var = len(train_X) input_dim = model.get_input_dim(train_X, num_var, curr_stride, curr_kernel_size) mdl = model.CnnLSTM(num_var, input_dim, output_dim, kernel_size=curr_kernel_size, stride=curr_stride, hidden_dim=curr_hidden_dim, num_lstm_layers=curr_num_layers, num_epochs=curr_num_epochs, learning_rate=curr_lr) mdl.to(device) history = mdl.fit_cv(train_loader, valid_X, valid_y, device) pred_y = mdl.predict(valid_X, device) history_all.append(history) test_rmse = np.sqrt(((valid_y - pred_y)**2).mean()) test_cos = np.asarray([ compute_cosine(valid_y[i, :], pred_y[i, :]) for i in range(len(valid_y)) ]).mean() score.append([test_rmse, test_cos]) cv_results = { 'score': score, 'parameter_all': parameter_all, 'history_all': history_all } save_results( rootpath + 'cv_results_test/cv_results_' + model_name + '_{}_{}.pkl'.format(cv_year, cv_index), cv_results)
def test(): imTransform = transforms.Compose([ torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) testLoader = helper.PreprocessInFocusOnlyData(imTransform, rootDirTest, 'inputs', 'masks', batchSize, False) sliceNames = ('outputSlice1', 'outputSlice10') numSlices = len(sliceNames) resultInputDir = rootDirResults + 'InputAIF/' resultDirInFocus = rootDirResults + 'OutputAIF/' resultDirSlices = rootDirResults + 'SliceOutputs/' resultDirSegMasks = rootDirResults + 'GroundTruthMasks/' extension = '.pth' loadRootDir = rootDirLoad epoch = loadEpochNum model1LoadPath = loadRootDir + 'model1/' model2LoadPath = loadRootDir + 'model2/' model3LoadPath = loadRootDir + 'model3/' print(model1LoadPath, model2LoadPath, model3LoadPath) model1 = model.EncoderDecoder(4, 6 * numSlices).to(device) model1.load_state_dict(torch.load(model1LoadPath + str(epoch) + extension)) model1.eval() model2 = model.DepthToInFocus(3).to(device) model2.load_state_dict(torch.load(model2LoadPath + str(epoch) + extension)) model2.eval() model3 = model.EncoderDecoderv2(3, 1).to(device) model3.load_state_dict(torch.load(model3LoadPath + str(epoch) + extension)) model3.eval() count = 0 for step, sample in enumerate(testLoader): print(step) inData = sample['Input'].to(device) print(inData.shape) segMask = sample['Mask'].to(device) blurOut = model1(inData, segMask) count += 1 for dispIndex in range(len(blurOut)): dispOut = blurOut[dispIndex] inFocus = F.interpolate(inData, size=None, scale_factor=1 / pow(2.0, dispIndex), mode='bilinear') blurMapSlices = helper.BlurMapsToSlices(dispOut, numSlices, 6) inFocusRegionSlices = [] nearFocusOutSlices = [] binaryBlurMapSlices = [] for slices in range(numSlices): nearFocusOutSlice = model2(inFocus, blurMapSlices[slices]) binaryBlurMapSlice = model3(nearFocusOutSlice) binaryBlurMapSlice = torch.cat( (binaryBlurMapSlice, binaryBlurMapSlice, binaryBlurMapSlice), dim=1) inFocusRegionSlice = torch.mul(nearFocusOutSlice, binaryBlurMapSlice) inFocusRegionSlices.append(inFocusRegionSlice) nearFocusOutSlices.append(nearFocusOutSlice) binaryBlurMapSlices.append(binaryBlurMapSlice) inFocusOutput = helper.CombineFocusRegionsToAIF( inFocusRegionSlices, numSlices, device) if dispIndex == 0: inFocusOutDispOut = helper.TensorToDispImage(inFocusOutput, 3) inDataDispOut = helper.TensorToDispImage(inFocus, 3) inSegMask = F.interpolate(segMask, size=None, scale_factor=1 / pow(2.0, dispIndex), mode='bilinear') save_image(inFocusOutDispOut, resultDirInFocus + str(step) + ".png", nrow=8, padding=2, normalize=False, range=None, scale_each=False, pad_value=0) save_image(inDataDispOut, resultInputDir + str(step) + ".png", nrow=8, padding=2, normalize=False, range=None, scale_each=False, pad_value=0) save_image(inSegMask, resultDirSegMasks + str(step) + ".png", nrow=8, padding=2, normalize=False, range=None, scale_each=False, pad_value=0) for indices in range(numSlices): dirToWrite = resultDirSlices + sliceNames[indices] + '/' sliceDir = dirToWrite + 'FocalSlices/' blurMapDir = dirToWrite + 'BlurMaps/' binaryMapDir = dirToWrite + 'BinaryMaps/' inFocusRegionDir = dirToWrite + 'InFocusRegion/' focalSliceDispOut = helper.TensorToDispImage( nearFocusOutSlices[indices], 3) binaryMapDispOut = binaryBlurMapSlices[indices] inFocusRegionDispOut = helper.TensorToDispImage( inFocusRegionSlices[indices], 3) displayOut = blurMapSlices[indices].mean(1) displayOutDispOut = displayOut.new(*displayOut.size()) displayOutDispOut[ 0, :, :] = displayOut[0, :, :] * 0.5 + 0.5 save_image(focalSliceDispOut, sliceDir + str(step) + ".png", nrow=8, padding=2, normalize=False, range=None, scale_each=False, pad_value=0) save_image(displayOutDispOut, blurMapDir + str(step) + ".png", nrow=8, padding=2, normalize=False, range=None, scale_each=False, pad_value=0) save_image(binaryMapDispOut, binaryMapDir + str(step) + ".png", nrow=8, padding=2, normalize=False, range=None, scale_each=False, pad_value=0) save_image(inFocusRegionDispOut, inFocusRegionDir + str(step) + ".png", nrow=8, padding=2, normalize=False, range=None, scale_each=False, pad_value=0)
def forecast_rep(month_id, year, rootpath, param_path, device, model_name, num_rep): """Run encoder-decoder style models with repetition - results are saved in a folder named forecast_results Args: month_id: an int indicating the month which is being forecasted year: an int indicating the year which is being forecasted rootpath: the path where the training and test sets are saved param_path: the path where the best hyperparameters are saved device: an indication if the model is runing on GPU or CPU model_name: a string indicating the name of a model num_rep: the number of repetition """ results = {} results['prediction_train'] = [] results['prediction_test'] = [] train_X = load_results( rootpath + 'train_X_pca_{}_forecast{}.pkl'.format(year, month_id)) test_X = load_results( rootpath + 'test_X_pca_{}_forecast{}.pkl'.format(year, month_id)) train_y = load_results( rootpath + 'train_y_pca_{}_forecast{}.pkl'.format(year, month_id)) test_y = load_results( rootpath + 'test_y_pca_{}_forecast{}.pkl'.format(year, month_id)) if model_name == 'EncoderFNN_AllSeq_AR_CI' or model_name == 'EncoderFNN_AllSeq_AR': train_y_ar = load_results( rootpath + 'train_y_pca_ar_{}_forecast{}.pkl'.format(year, month_id)) test_y_ar = load_results( rootpath + 'test_y_pca_ar_{}_forecast{}.pkl'.format(year, month_id)) input_dim = train_X.shape[-1] output_dim = train_y.shape[-1] # ar_dim = train_X.shape[1] best_parameter = load_results( param_path + '{}_forecast{}.pkl'.format(model_name, month_id)) for rep in range(num_rep): if model_name == 'EncoderDecoder': curr_hidden_dim = best_parameter['hidden_dim'] curr_num_layer = best_parameter['num_layers'] curr_decoder_len = best_parameter['decoder_len'] curr_threshold = best_parameter['threshold'] curr_lr = best_parameter['learning_rate'] curr_num_epochs = best_parameter['num_epochs'] mdl = model.EncoderDecoder(input_dim=input_dim, output_dim=output_dim, hidden_dim=curr_hidden_dim, num_layers=curr_num_layer, learning_rate=curr_lr, decoder_len=curr_decoder_len, threshold=curr_threshold, num_epochs=curr_num_epochs) # set data for training train_dataset = model.MapDataset(train_X, train_y) train_loader = DataLoader(dataset=train_dataset, batch_size=512, shuffle=False) model.init_weight(mdl) # send model to gpu mdl.to(device) mdl.fit(train_loader, device) state = {'state_dict': mdl.state_dict()} torch.save( state, rootpath + 'models/{}_{}_{}.t7'.format(model_name, year, month_id)) pred_train = mdl.predict(train_X, device) pred_y = mdl.predict(test_X, device) elif model_name == 'EncoderFNN': curr_hidden_dim = best_parameter['hidden_dim'] curr_num_layer = best_parameter['num_layers'] curr_seq_len = best_parameter['seq_len'] curr_threshold = best_parameter['threshold'] curr_lr = best_parameter['learning_rate'] curr_num_epochs = best_parameter['num_epochs'] curr_last_layer = best_parameter['last_layer'] mdl = model.EncoderFNN(input_dim=input_dim, output_dim=output_dim, hidden_dim=curr_hidden_dim, num_layers=curr_num_layer, last_layer=curr_last_layer, seq_len=curr_seq_len, learning_rate=curr_lr, threshold=curr_threshold, num_epochs=curr_num_epochs) # set data for training train_dataset = model.MapDataset(train_X, train_y) train_loader = DataLoader(dataset=train_dataset, batch_size=512, shuffle=False) model.init_weight(mdl) # send model to gpu mdl.to(device) mdl.fit(train_loader, device) state = {'state_dict': mdl.state_dict()} torch.save( state, rootpath + 'models/{}_{}_{}.t7'.format(model_name, year, month_id)) pred_train = mdl.predict(train_X, device) pred_y = mdl.predict(test_X, device) elif model_name == 'EncoderFNN_AllSeq': curr_hidden_dim = best_parameter['hidden_dim'] curr_num_layer = best_parameter['num_layers'] curr_seq_len = best_parameter['seq_len'] curr_threshold = best_parameter['threshold'] curr_lr = best_parameter['learning_rate'] curr_num_epochs = best_parameter['num_epochs'] curr_linear_dim = best_parameter['linear_dim'] curr_drop_out = best_parameter['drop_out'] mdl = model.EncoderFNN_AllSeq(input_dim=input_dim, output_dim=output_dim, hidden_dim=curr_hidden_dim, num_layers=curr_num_layer, seq_len=curr_seq_len, linear_dim=curr_linear_dim, learning_rate=curr_lr, dropout=curr_drop_out, threshold=curr_threshold, num_epochs=curr_num_epochs) # set data for training train_dataset = model.MapDataset(train_X, train_y) train_loader = DataLoader(dataset=train_dataset, batch_size=512, shuffle=False) model.init_weight(mdl) # send model to gpu mdl.to(device) mdl.fit(train_loader, device) state = {'state_dict': mdl.state_dict()} torch.save( state, rootpath + 'models/{}_{}_{}.t7'.format(model_name, year, month_id)) pred_train = mdl.predict(train_X, device) pred_y = mdl.predict(test_X, device) elif model_name == 'EncoderFNN_AllSeq_AR': curr_hidden_dim = best_parameter['hidden_dim'] curr_num_layer = best_parameter['num_layers'] curr_seq_len = best_parameter['seq_len'] curr_threshold = best_parameter['threshold'] curr_lr = best_parameter['learning_rate'] curr_num_epochs = best_parameter['num_epochs'] curr_linear_dim = best_parameter['linear_dim'] curr_drop_out = best_parameter['drop_out'] mdl = model.EncoderFNN_AllSeq_AR(input_dim=input_dim, output_dim=output_dim, hidden_dim=curr_hidden_dim, num_layers=curr_num_layer, seq_len=curr_seq_len, linear_dim=curr_linear_dim, learning_rate=curr_lr, dropout=curr_drop_out, threshold=curr_threshold, num_epochs=curr_num_epochs) train_dataset = model.MapDataset_ar(train_X, train_y_ar, train_y) train_loader = DataLoader(dataset=train_dataset, batch_size=512, shuffle=False) model.init_weight(mdl) # send model to gpu mdl.to(device) mdl.fit(train_loader, device) state = {'state_dict': mdl.state_dict()} torch.save( state, rootpath + 'models/{}_{}_{}.t7'.format(model_name, year, month_id)) pred_train = mdl.predict(train_X, train_y_ar, device) pred_y = mdl.predict(test_X, test_y_ar, device) elif model_name == 'EncoderFNN_AllSeq_AR_CI': curr_hidden_dim = best_parameter['hidden_dim'] curr_num_layer = best_parameter['num_layers'] curr_seq_len = best_parameter['seq_len'] curr_threshold = best_parameter['threshold'] curr_lr = best_parameter['learning_rate'] curr_num_epochs = best_parameter['num_epochs'] curr_linear_dim = best_parameter['linear_dim'] curr_drop_out = best_parameter['drop_out'] ci_dim = best_parameter['ci_dim'] mdl = model.EncoderFNN_AllSeq_AR_CI(input_dim=input_dim - ci_dim, output_dim=output_dim, hidden_dim=curr_hidden_dim, num_layers=curr_num_layer, seq_len=curr_seq_len, linear_dim=curr_linear_dim, ci_dim=ci_dim, learning_rate=curr_lr, dropout=curr_drop_out, threshold=curr_threshold, num_epochs=curr_num_epochs) train_dataset = model.MapDataset_ar(train_X, train_y_ar, train_y) train_loader = DataLoader(dataset=train_dataset, batch_size=512, shuffle=False) model.init_weight(mdl) # send model to gpu mdl.to(device) mdl.fit(train_loader, device) state = {'state_dict': mdl.state_dict()} torch.save( state, rootpath + 'models/{}_{}_{}.t7'.format(model_name, year, month_id)) pred_train = mdl.predict(train_X, train_y_ar, device) pred_y = mdl.predict(test_X, test_y_ar, device) results['target_train'] = train_y results['prediction_train'].append(pred_train) results['target_test'] = test_y results['prediction_test'].append(pred_y) save_results( rootpath + 'forecast_results/results_{}_{}_{}.pkl'.format( model_name, year, month_id), results)
def ids2words(lang, ids): return [lang.index2word[idx] for idx in ids] def greedy_decode(model, dataloader, input_lang, output_lang): with torch.no_grad(): batch = next(iter(dataloader)) input_tensor = batch[0] input_mask = batch[1] target_tensor = batch[2] decoder_outputs, decoder_hidden = model(input_tensor, input_mask) topv, topi = decoder_outputs.topk(1) decoded_ids = topi.squeeze() for idx in range(input_tensor.size(0)): input_sent = ids2words(input_lang, input_tensor[idx].cpu().numpy()) output_sent = ids2words(output_lang, decoded_ids[idx].cpu().numpy()) target_sent = ids2words(output_lang, target_tensor[idx].cpu().numpy()) print('Input: {}'.format(input_sent)) print('Target: {}'.format(target_sent)) print('Output: {}'.format(output_sent)) if __name__ == '__main__': input_lang, output_lang, train_dataloader = load_data.get_dataloader(batch_size) model = model.EncoderDecoder(hidden_size, input_lang.n_words, output_lang.n_words).to(device) train(train_dataloader, model, n_epochs=20) greedy_decode(model, train_dataloader, input_lang, output_lang)
def train(epochs): imTransform = transforms.Compose([ torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) trainLoader = helper.PreprocessInFocusOnlyData(imTransform, rootDirTrain, batchSize, True) validLoader = helper.PreprocessInFocusOnlyData(imTransform, rootDirValid, batchSize, True) slices = ('outputSlice1', 'outputSlice4', 'outputSlice7', 'outputSlice10') numSlices = len(slices) rootDirToSave = rootDirSave model1SaveDir = rootDirToSave + 'model1/' model2SaveDir = rootDirToSave + 'model2/' model3SaveDir = rootDirToSave + 'model3/' extension = '.pth' loadRootDir = rootDirSave model1LoadPath = loadRootDir + 'model1/' model2LoadPath = loadRootDir + 'model2/' model3LoadPath = loadRootDir + 'model3/' model1 = model.EncoderDecoder(3, 6 * numSlices).to(device) model2 = model.DepthToInFocus(3).to(device) model3 = model.EncoderDecoderv2(3, 1).to(device) model1.load_state_dict( torch.load(model1LoadPath + str(loadEpochNum) + extension)) model2.load_state_dict( torch.load(model2LoadPath + str(loadEpochNum) + extension)) model3.load_state_dict( torch.load(model3LoadPath + str(loadEpochNum) + extension)) contentLoss = losses.PerceptualLoss() contentLoss.initialize(nn.MSELoss()) trainLossArray = np.zeros(epochs) validLossArray = np.zeros(epochs) lr = learningRate minValidEpoch = 0 minValidLoss = 100.0 for epoch in range(epochs): start = time.clock() epochLoss = np.zeros(2) epochCount = epoch + 1 print("Epoch num : " + str(epochCount)) if epoch >= 20 and epoch % 20 == 0: lr = lr / 2 print('Learning rate changed to ' + str(lr)) for step, sample in enumerate(trainLoader): loss = torch.zeros(1).to(device) loss.requires_grad = False data = sample.to(device) blurOut = model1(data) for dispIndex in range(len(blurOut)): scaleLoss = torch.zeros(1).to(device) scaleLoss.requires_grad = False dispOut = blurOut[dispIndex] inFocus = F.interpolate(data, size=None, scale_factor=1 / pow(2.0, dispIndex), mode='bilinear') blurMapSlices = helper.BlurMapsToSlices(dispOut, numSlices, 6) inFocusRegionSlices = [] for slices in range(numSlices): nearFocusOutSlice = model2(inFocus, blurMapSlices[slices]) binaryBlurMapSlice = model3(nearFocusOutSlice) inFocusRegionSlice = torch.mul(nearFocusOutSlice, binaryBlurMapSlice) inFocusRegionSlices.append(inFocusRegionSlice) inFocusOutput = helper.CombineFocusRegionsToAIF( inFocusRegionSlices, numSlices, device) perceptLoss = contentLoss.get_loss(inFocus, inFocusOutput) simiLoss = losses.similarityLoss(inFocus, inFocusOutput, alpha, winSize) scaleLoss += perceptLossW * perceptLoss + simiLossW * simiLoss loss += scaleLoss epochLoss[0] += loss optimizer = optim.Adam(list(model1.parameters()) + list(model2.parameters()) + list(model3.parameters()), lr=lr,\ betas=(beta1, beta2), eps=1e-08, weight_decay=0, amsgrad=False) optimizer.zero_grad() loss.backward() optimizer.step() if step % 500 == 0: print("Step num :" + str(step) + " Train Loss :", str(loss.item())) with torch.no_grad(): for validStep, validSample in enumerate(validLoader): validLoss = torch.zeros(1).to(device) validLoss.requires_grad = False validData = sample.to(device) validBlurOut = model1(validData) for validDispIndex in range(len(validBlurOut)): validScaleLoss = torch.zeros(1).to(device) validScaleLoss.requires_grad = False validDispOut = validBlurOut[validDispIndex] validInFocus = F.interpolate(validData, size=None, scale_factor=1 / pow(2.0, validDispIndex), mode='bilinear') validBlurMapSlices = helper.BlurMapsToSlices( validDispOut, numSlices, 6) validInFocusRegionSlices = [] for validSlices in range(numSlices): validNearFocusOutSlice = model2( validInFocus, validBlurMapSlices[validSlices]) validBinaryBlurMapSlice = model3( validNearFocusOutSlice) validInFocusRegionSlice = torch.mul( validNearFocusOutSlice, validBinaryBlurMapSlice) validInFocusRegionSlices.append( validInFocusRegionSlice) validInFocusOutput = helper.CombineFocusRegionsToAIF( validInFocusRegionSlices, numSlices, device) validPerceptLoss = contentLoss.get_loss( validInFocus, validInFocusOutput) validSimiLoss = losses.similarityLoss( validInFocus, validInFocusOutput, alpha, winSize) validScaleLoss += perceptLossW * validPerceptLoss + simiLossW * validSimiLoss validLoss += validScaleLoss epochLoss[1] += validLoss epochLoss[0] /= step epochLoss[1] /= validStep epochFreq = epochSaveFreq print("Time taken for epoch " + str(epoch) + " is " + str(time.clock() - start)) if validLoss < minValidLoss: minValidLoss = validLoss minValidEpoch = epoch torch.save(model1.state_dict(), model1SaveDir + str(epoch) + 'generic_' + extension) torch.save(model2.state_dict(), model2SaveDir + str(epoch) + 'generic_' + extension) torch.save(model3.state_dict(), model3SaveDir + str(epoch) + 'generic_' + extension) print("Saving latest model at epoch: " + str(epoch)) if epochCount % epochFreq == 0: torch.save(model1.state_dict(), model1SaveDir + str(epoch) + 'generic_' + extension) torch.save(model2.state_dict(), model2SaveDir + str(epoch) + 'generic_' + extension) torch.save(model3.state_dict(), model3SaveDir + str(epoch) + 'generic_' + extension) print("Training loss for epoch " + str(epochCount) + " : " + str(epochLoss[0])) print("Validation loss for epoch " + str(epochCount) + " : " + str(epochLoss[1])) trainLossArray[epoch] = epochLoss[0] validLossArray[epoch] = epochLoss[1] print(minValidEpoch) return (trainLossArray, validLossArray)
def train(epochs) : imTransform = transforms.Compose([torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) sliceNames = ('outputSlice1','outputSlice10') numSlices = len(sliceNames) trainLoaderLF = helper.PreprocessData(imTransform, sliceNames, rootDirTrainLF, 'input', 'segments', batchSize, True) validLoaderLF = helper.PreprocessData(imTransform, sliceNames, rootDirValidLF, 'input', 'segments', batchSize, True) trainLoaderReal = helper.PreprocessInFocusOnlyData(imTransform, rootDirTrainReal, 'inputs', 'masks', batchSize, False) validLoaderReal = helper.PreprocessInFocusOnlyData(imTransform, rootDirValidReal, 'inputs', 'masks', batchSize, False) rootDirToSave = rootDirSave model1SaveDir = rootDirToSave + 'model1/' model2SaveDir = rootDirToSave + 'model2/' model3SaveDir = rootDirToSave + 'model3/' extension = '.pth' model1 = model.EncoderDecoder(4, 6 * numSlices).to(device) model2 = model.DepthToInFocus(3).to(device) model3 = model.EncoderDecoderv2(3, 1).to(device) contentLoss = losses.PerceptualLoss() contentLoss.initialize(nn.MSELoss()) trainLossArray = np.zeros(epochs) validLossArray = np.zeros(epochs) lr = learningRate minValidEpoch = 0 minValidLoss = 100.0 trainLoaderRealList = list(trainLoaderReal) validLoaderRealList = list(validLoaderReal) numRealSteps = len(trainLoaderRealList) for epoch in range(epochs) : start = time.clock() realStep = 0 epochLoss = np.zeros(2) epochCount = epoch + 1 print ("Epoch num : " + str(epochCount)) if epoch >= 20 and epoch % 20 == 0 : lr = lr / 2 print ('Learning rate changed to ' + str(lr)) if epoch < 40 : binarizationW = 0.0 for step, sample in enumerate(trainLoaderLF) : if step % 3 == 1 and realStep < numRealSteps: realStepLoss = torch.zeros(1).to(device) realInData = trainLoaderRealList[realStep]['Input'].to(device) realSegMask = trainLoaderRealList[realStep]['Mask'].to(device) realBlurOut = model1(realInData, realSegMask) for realDispIndex in range(len(realBlurOut)) : realScaleLoss = torch.zeros(1).to(device) realBinarizationLoss = torch.zeros(1).to(device) realDispOut = realBlurOut[realDispIndex] realInFocus = F.interpolate(realInData, size = None, scale_factor = 1/pow(2.0, realDispIndex), mode = 'bilinear') realBlurMapSlices = helper.BlurMapsToSlices(realDispOut, numSlices, 6) realInFocusRegionSlices = [] for realSlices in range(numSlices) : realNearFocusOutSlice = model2(realInFocus, realBlurMapSlices[realSlices]) realBinaryBlurMapSlice = model3(realNearFocusOutSlice) realBinaryBlurMapSlice = torch.cat((realBinaryBlurMapSlice, realBinaryBlurMapSlice, realBinaryBlurMapSlice), dim = 1) binaryLoss = losses.binarizationLoss(realBinaryBlurMapSlice) * binarizationLossW realBinarizationLoss += binaryLoss realInFocusRegionSlice = torch.mul(realNearFocusOutSlice, realBinaryBlurMapSlice) realInFocusRegionSlices.append(realInFocusRegionSlice) realInFocusOutput = helper.CombineFocusRegionsToAIF(realInFocusRegionSlices, numSlices, device) realPerceptLoss = contentLoss.get_loss(realInFocus, realInFocusOutput) realSimiLoss = losses.similarityLoss(realInFocus, realInFocusOutput, alpha, winSize) realScaleLoss += perceptLossW * realPerceptLoss + simiLossW * realSimiLoss + realBinarizationLoss realStepLoss += realScaleLoss loss += realStepLoss optimizer = optim.Adam(list(model1.parameters()) + list(model2.parameters()) + list(model3.parameters()), lr=lr/5,\ betas=(beta1, beta2), eps=1e-08, weight_decay=0, amsgrad=False) optimizer.zero_grad() realStepLoss.backward(retain_graph = True) optimizer.step() realStep += 1 stepLossStage1 = torch.zeros(1).to(device) stepLossStage1.requires_grad = False stepLossStage2 = torch.zeros(1).to(device) stepLossStage2.requires_grad = False loss = torch.zeros(1).to(device) loss.requires_grad = False inData = sample['Input'].to(device) outData = sample['Output'].to(device) inSegMask = sample['Mask'].to(device) blurOut = model1(inData, inSegMask) for dispIndex in range(len(blurOut)) : scaleLossStage1 = torch.zeros(1).to(device) scaleLossStage1.requires_grad = False scaleLossStage2 = torch.zeros(1).to(device) scaleLossStage2.requires_grad = False dispOut = blurOut[dispIndex] inFocus = F.interpolate(inData, size = None, scale_factor = 1 / pow(2.0, dispIndex), mode = 'bilinear') nearFocusStackGT = F.interpolate(outData, size = None, scale_factor = 1 / pow(2.0, dispIndex) , mode = 'bilinear') nearFocusGTSlices = helper.FocalStackToSlices(nearFocusStackGT, numSlices) blurMapSlices = helper.BlurMapsToSlices(dispOut, numSlices, 6) inFocusRegionSlices = [] for slices in range(numSlices) : sliceLoss = torch.zeros(1).to(device) sliceLoss.requires_grad = False #print (time.clock() - start) nearFocusOutSlice = model2(inFocus, blurMapSlices[slices]) #print (time.clock() - start) perceptLoss = contentLoss.get_loss(nearFocusGTSlices[slices], nearFocusOutSlice) simiLoss = losses.similarityLoss(nearFocusGTSlices[slices], nearFocusOutSlice, alpha, winSize) sliceLoss = perceptLossW * perceptLoss + simiLoss * simiLoss binaryBlurMapSlice = model3(nearFocusOutSlice) binaryBlurMapSlice = torch.cat((binaryBlurMapSlice, binaryBlurMapSlice, binaryBlurMapSlice), dim = 1) binaryLoss = losses.binarizationLoss(binaryBlurMapSlice) * binarizationLossW sliceLoss += binaryLoss inFocusRegionSlice = torch.mul(nearFocusOutSlice, binaryBlurMapSlice) inFocusRegionSlices.append(inFocusRegionSlice) scaleLossStage1 += sliceLoss inFocusOutput = helper.CombineFocusRegionsToAIF(inFocusRegionSlices, numSlices, device) perceptLossStage2 = contentLoss.get_loss(inFocus, inFocusOutput) simiLossStage2 = losses.similarityLoss(inFocus, inFocusOutput, alpha, winSize) scaleLossStage2 += perceptLossW * perceptLossStage2 + simiLossW * simiLossStage2 stepLossStage1 += scaleLossStage1 stepLossStage2 += scaleLossStage2 loss += stepLossStage1 + stepLossStage2 epochLoss[0] += loss if jointTraining == True : optimizer = optim.Adam(list(model1.parameters()) + list(model2.parameters()) + list(model3.parameters()), lr=lr,\ betas=(beta1, beta2), eps=1e-08, weight_decay=0, amsgrad=False) optimizer.zero_grad() loss.backward() optimizer.step() else : stage1Parameters = list(model1.parameters()) + list(model2.parameters()) stage2Parameters = list(model3.parameters()) optimizerStage1 = optim.Adam(stage1Parameters, lr=lr, betas=(beta1, beta2), eps=1e-08, weight_decay=0, amsgrad=False) optimizerStage2 = optim.Adam(stage2Parameters, lr=lr, betas=(beta1, beta2), eps=1e-08, weight_decay=0, amsgrad=False) optimizerStage1.zero_grad() stepLossStage1.backward(retain_graph = True) optimizerStage1.step() optimizerStage2.zero_grad() stepLossStage2.backward() optimizerStage2.step() if step % 500 == 0 : print ("Step num :" + str(step) + " Train Loss :", str(loss.item())) with torch.no_grad() : validRealStep = 0 for validStep, validSample in enumerate(validLoaderLF) : if validRealStep % 3 == 1 and validRealStep < numRealSteps: validRealStepLoss = torch.zeros(1).to(device) validRealInData = validLoaderRealList[realStep]['Input'].to(device) validRealSegMask = validLoaderRealList[realStep]['Mask'].to(device) validRealBlurOut = model1(validRealInData, validRealSegMask) for validRealDispIndex in range(len(validRealBlurOut)) : validRealScaleLoss = torch.zeros(1).to(device) validRealDispOut = validRealBlurOut[validRealDispIndex] validRealInFocus = F.interpolate(validRealInData, size = None, scale_factor = 1/pow(2.0, validRealDispIndex), mode = 'bilinear') validRealBlurMapSlices = helper.BlurMapsToSlices(validRealDispOut, numSlices, 6) validRealInFocusRegionSlices = [] for validRealSlices in range(numSlices) : validRealNearFocusOutSlice = model2(validRealInFocus, validRealBlurMapSlices[validRealSlices]) validRealBinaryBlurMapSlice = model3(validRealNearFocusOutSlice) validRealBinaryBlurMapSlice = torch.cat((validRealBinaryBlurMapSlice, validRealBinaryBlurMapSlice, validRealBinaryBlurMapSlice), dim = 1) validRealInFocusRegionSlice = torch.mul(validRealNearFocusOutSlice, validRealBinaryBlurMapSlice) validRealInFocusRegionSlices.append(validRealInFocusRegionSlice) validRealInFocusOutput = helper.CombineFocusRegionsToAIF(validRealInFocusRegionSlices, numSlices, device) validRealPerceptLoss = contentLoss.get_loss(validRealInFocus, validRealInFocusOutput) validRealSimiLoss = losses.similarityLoss(validRealInFocus, validRealInFocusOutput, alpha, winSize) validRealScaleLoss += perceptLossW * validRealPerceptLoss + simiLossW * validRealSimiLoss validRealStepLoss += validRealScaleLoss validRealStep += 1 validLoss += validRealStepLoss validStepLossStage1 = torch.zeros(1).to(device) validStepLossStage1.requires_grad = False validStepLossStage2 = torch.zeros(1).to(device) validStepLossStage2.requires_grad = False validLoss = torch.zeros(1).to(device) validLoss.requires_grad = False validInData = validSample['Input'].to(device) validOutData = validSample['Output'].to(device) validInSegMask = validSample['Mask'].to(device) validBlurOut = model1(validInData, validInSegMask) for validDispIndex in range(len(validBlurOut)) : validScaleLossStage1 = torch.zeros(1).to(device) validScaleLossStage1.requires_grad = False validScaleLossStage2 = torch.zeros(1).to(device) validScaleLossStage2.requires_grad = False validDispOut = validBlurOut[validDispIndex] validInFocus = F.interpolate(validInData, size = None, scale_factor = 1 / pow(2.0, validDispIndex), mode = 'bilinear') validNearFocusStackGT = F.interpolate(validOutData, size = None, scale_factor = 1 / pow(2.0, validDispIndex) , mode = 'bilinear') validNearFocusGTSlices = helper.FocalStackToSlices(validNearFocusStackGT, numSlices) validBlurMapSlices = helper.BlurMapsToSlices(validDispOut, numSlices, 6) validInFocusRegionSlices = [] for validSlices in range(numSlices) : validSliceLoss = torch.zeros(1).to(device) validSliceLoss.requires_grad = False validNearFocusOutSlice = model2(validInFocus, validBlurMapSlices[validSlices]) validPerceptLoss = contentLoss.get_loss(validNearFocusGTSlices[slices], validNearFocusOutSlice) validSimiLoss = losses.similarityLoss(validNearFocusGTSlices[slices], validNearFocusOutSlice, alpha, winSize) validSliceLoss = perceptLossW * validPerceptLoss + validSimiLoss * validSimiLoss validBinaryBlurMapSlice = model3(validNearFocusOutSlice) validBinaryBlurMapSlice = torch.cat((validBinaryBlurMapSlice, validBinaryBlurMapSlice, validBinaryBlurMapSlice), dim = 1) validBinaryLoss = losses.binarizationLoss(validBinaryBlurMapSlice) * binarizationLossW validSliceLoss += validBinaryLoss validInFocusRegionSlice = torch.mul(validNearFocusOutSlice, validBinaryBlurMapSlice) validInFocusRegionSlices.append(validInFocusRegionSlice) validScaleLossStage1 += validSliceLoss validInFocusOutput = helper.CombineFocusRegionsToAIF(validInFocusRegionSlices, numSlices, device) validPerceptLossStage2 = contentLoss.get_loss(validInFocus, validInFocusOutput) validSimiLossStage2 = losses.similarityLoss(validInFocus, validInFocusOutput, alpha, winSize) validScaleLossStage2 += perceptLossW * validPerceptLossStage2 + simiLossW * validSimiLossStage2 validStepLossStage1 += validScaleLossStage1 validStepLossStage2 += validScaleLossStage2 validLoss += validStepLossStage1 + validStepLossStage2 epochLoss[1] += validLoss epochLoss[0] /= step epochLoss[1] /= validStep epochFreq = epochSaveFreq print ("Time taken for epoch " + str(epoch) + " is " + str(time.clock() - start)) if validLoss < minValidLoss : minValidLoss = validLoss minValidEpoch = epoch torch.save(model1.state_dict(), model1SaveDir + str(epoch) + extension) torch.save(model2.state_dict(), model2SaveDir + str(epoch) + extension) torch.save(model3.state_dict(), model3SaveDir + str(epoch) + extension) print ("Saving latest model at epoch: " + str(epoch)) if epochCount % epochFreq == 0 : torch.save(model1.state_dict(), model1SaveDir + str(epoch) + extension) torch.save(model2.state_dict(), model2SaveDir + str(epoch) + extension) torch.save(model3.state_dict(), model3SaveDir + str(epoch) + extension) print ("Training loss for epoch " + str(epochCount) + " : " + str(epochLoss[0])) print ("Validation loss for epoch " + str(epochCount) + " : " + str(epochLoss[1])) trainLossArray[epoch] = epochLoss[0] validLossArray[epoch] = epochLoss[1] print (minValidEpoch) return (trainLossArray, validLossArray)