def train_datagen(epoch_iter=2000, epoch_num=5, batch_size=128, data_dir=args.train_data): print('train_datagen') while (True): n_count = 0 if n_count == 0: print(n_count) xs = dg.datagenerator(data_dir) print(xs.shape) #assert len(xs)%args.batch_size ==0, \ print('log done') xs = xs / 255.0 xs = xs.astype('float32') indices = list(range(xs.shape[0])) n_count = 1 print('if done', range(epoch_num)) for _ in range(epoch_num): np.random.shuffle(indices) # shuffle for i in range(0, len(indices), batch_size): batch_x = xs[indices[i:i + batch_size]] noise = np.random.normal(0, args.sigma / 255.0, batch_x.shape) # noise #noise = K.random_normal(ge_batch_y.shape, mean=0, stddev=args.sigma/255.0) batch_y = batch_x + noise yield batch_y, batch_x
def train_datagen(epoch_iter=1, epoch_num=1, batch_size=128, data_dir=args.train_data, noise_dir=args.noise_data): loop = 0 while (True): loop = loop + 1 xs = dg.datagenerator(data_dir) assert len(xs)%batch_size ==0, \ log('make sure the last iteration has a full batchsize, this is important if you use batch normalization!') xs = xs.astype('float32') indices = list(range(xs.shape[0])) print("Augmented patch data shape:", xs.shape) np.random.shuffle(indices) # shuffle for i in range(0, len(indices), batch_size): ratio1 = np.random.uniform(0.7, 0.3, batch_size) ratio2 = 1.0 - ratio1 batch_x = xs[indices[i:i + batch_size]] noise = extract_noise(batch_x.shape, noise_dir) for jj in range(batch_size): batch_x[jj] = batch_x[jj] * ratio1[jj] noise[jj] = noise[jj] * ratio2[jj] batch_y = (batch_x + noise) yield batch_y, batch_x
def main(): global opt, model opt = parser.parse_args() logger = set_logger(opt.save) print(opt) print(opt, file=logger) # setting gpu and seed os.environ["CUDA_VISIBLE_DEVICES"] = '0' opt.seed = random.randint(1, 10000) print("Random Seed: ", opt.seed) torch.manual_seed(opt.seed) torch.cuda.manual_seed(opt.seed) # setting dataset print("===> Loading dataset") patches = datagenerator(data_dir=opt.data_train) train_set = DenoisingDataset(patches, sigma=opt.sigma) training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, drop_last=True, batch_size=opt.batchSize, shuffle=True) # setting model and loss print("===> Building model") model = DUAL_CNN_DENOISE() criterion = nn.MSELoss(size_average=False) model = model.cuda() criterion = criterion.cuda() # setting optimizer print("===> Setting Optimizer") kwargs = {'weight_decay': opt.weight_decay} optimizer = optim.Adam([{ "params": model.structure_net.parameters(), "lr": opt.srcnn_lr }, { "params": model.detail_net.parameters(), "lr": opt.vdsr_lr }], **kwargs) print("===> Training") for epoch in range(opt.start_epoch, opt.nEpochs + 1): train(training_data_loader, optimizer, model, criterion, epoch, logger) model_path = save_checkpoint(model, epoch) eval.eval(model_path, opt.save, opt.sigma)
def train_datagen(epoch_iter=2000,epoch_num=5,batch_size=128,data_dir=args.train_data): while(True): n_count = 0 if n_count == 0: #print(n_count) xs, ys = dg.datagenerator(data_dir) assert len(xs)%args.batch_size ==0, \ log('make sure the last iteration has a full batchsize, this is important if you use batch normalization!') xs = xs.astype('float32')/255.0 ys = ys.astype('float32')/255.0 indices = list(range(xs.shape[0])) n_count = 1 for _ in range(epoch_num): np.random.shuffle(indices) # shuffle for i in range(0, len(indices), batch_size): batch_x = xs[indices[i:i+batch_size]] batch_y = ys[indices[i:i+batch_size]] yield batch_y, batch_x
def train(self): train_images, train_labels = datagenerator(train_data_dir=self.train_images_dir, GT_data_dir=self.train_labels_dir) train_images = train_images / 255 train_labels = train_labels / 255 self.summary_writer = tf.summary.FileWriter(self.log_dir, graph=tf.get_default_graph()) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): self.train_op = tf.train.AdamOptimizer(self.lr_init).minimize(self.loss) tf.initialize_all_variables().run() counter = 0 start_time = time.time() if self.load(self.checkpoint_dir): print(" [*] Load SUCCESS") else: print(" [!] Load failed...") print("Training...") for ep in range(self.epoch): if ep % 5 ==0 and ep != 0: self.lr_init = self.lr_init / 10 batch_idxs = len(train_images) // self.batch_size for idx in range(0, batch_idxs): batch_images = train_images[idx * self.batch_size: (idx + 1) * self.batch_size] batch_labels = train_labels[idx * self.batch_size: (idx + 1) * self.batch_size] counter += 1 _, err, psnr = self.sess.run([self.train_op, self.loss, self.metric], feed_dict={self.images:batch_images, self.labels:batch_labels}) summary = self.sess.run(self.merged_summary_op, feed_dict={self.images:batch_images, self.labels:batch_labels}) self.summary_writer.add_summary(summary, counter) if counter % 10 == 0: print("Epoch: [%2d], step: [%2d], time: [%4.4f], loss: [%.8f], PSNR: [%.4f], lr: [%.6f]" % ( (ep + 1), counter, time.time() - start_time, err, psnr, self.lr_init)) if counter % 500 == 0: self.save(self.checkpoint_dir, counter)
optimizer_decompose = optim.Adam(decompose_model.parameters(), lr=args.lr) scheduler_decompose = MultiStepLR(optimizer_decompose, milestones=[30, 60, 90], gamma=0.2) # learning rates # optimizer_compose = optim.Adam(compose_model.parameters(), lr=args.lr) # scheduler_compose = MultiStepLR(optimizer_compose, milestones=[30, 60, 90], gamma=0.2) # learning rates for epoch in range(initial_epoch, n_epoch): decompose_model.train() # compose_model.train() scheduler_decompose.step( epoch) # step to the learning rate in this epcoh # scheduler_compose.step(epoch) # step to the learning rate in this epcoh xs = dg.datagenerator(data_dir=args.train_data) xs = xs.astype('float32') / 255.0 xs = torch.from_numpy(xs.transpose( (0, 3, 1, 2))) # tensor of the clean patches, NXCXHXW DDataset = DenoisingDataset(xs, sigma) batch_y, batch_x = DDataset[:238336] dataset = torch.cat((batch_x, batch_y), dim=1) DLoader = torch.utils.data.DataLoader(dataset=dataset, num_workers=0, drop_last=True, batch_size=batch_size, shuffle=True) epoch_loss = 0 start_time = time.time()
# DADOS PARA VALIDACAO DO TREINAMENTO # val_xs, val_illums = dg.datagenerator(data_dir=args.val_data, batch_size=args.batch_size, is_validation=True) # val_xs = val_xs.astype(np.float32)/255.0 #float = [0,1] # val_xs = torch.from_numpy(np.reshape(val_xs,(val_xs.shape[0], val_xs.shape[3], val_xs.shape[1], val_xs.shape[2]))) # print('val_xs',val_xs.shape) # val_DDataset = DenoisingDataset(val_xs, val_illums) # val_DLoader = DataLoader(dataset=val_DDataset, num_workers=4, drop_last=True, batch_size=batch_size, shuffle=True, # worker_init_fn=worker_init) for epoch in range(initial_epoch, n_epoch): print('Treinando epoca %d' % (epoch)) # print('Treinando epoca %d com fator %f'%(epoch, fatores[epoch])) # scheduler.step(epoch) xs, illums = dg.datagenerator(data_dir=args.train_data, batch_size=args.batch_size) xs = xs.astype(np.float32) / 255.0 # float = [0,1] xs = torch.from_numpy(np.reshape(xs, ( xs.shape[0], xs.shape[3], xs.shape[1], xs.shape[2]))) # tensor of the clean patches, NXCXHXW # xs = torch.from_numpy(xs.transpose((0, 3, 1, 2))) # tensor of the clean patches, NXCXHXW # xs = xs.transpose((0, 3, 1, 2)) DDataset = LightingDataset(xs, illums) DLoader = DataLoader(dataset=DDataset, num_workers=4, drop_last=True, batch_size=batch_size, shuffle=True, worker_init_fn=worker_init) epoch_loss = 0
def train(batch_size=128, n_epoch=300, sigma=25, lr=1e-4, depth=7, device="cuda:0", data_dir='./data/Train400', model_dir='models', model_name=None): device = torch.device(device) if not os.path.exists( os.path.join( model_dir, "model" + str(sigma) + "m" + str(model_name[1]) + "d" + str(depth))): os.mkdir( os.path.join( model_dir, "model" + str(sigma) + "m" + str(model_name[1]) + "d" + str(depth))) save_dir = os.path.join( model_dir, "model" + str(sigma) + "m" + str(model_name[1]) + "d" + str(depth)) from datetime import date save_name = "model_mode" + str(model_name[1]) + str(depth) + "_" + "".join( str(date.today()).split('-')[1:]) + ".pth" f = open(os.path.join(save_dir, save_name.replace(".pth", ".txt")), 'w') f.write(('--\t This is end to end model saved as ' + save_name + '\n')) f.write(('--\t epoch %4d batch_size %4d sigma %4d\n' % (n_epoch, batch_size, sigma))) f.write(model_name[0]) DNet = torch.load(os.path.join(model_dir, model_name[0])) DNet.eval() modelG = Model(model_dir=model_dir, model_name=model_name) #guidance='noisy, denoised' modelD = discriminator() print(modelG) f.write(str(modelG)) f.write('\n\n') ngpu = 2 if (device.type == 'cuda') and (ngpu > 1): modelG = nn.DataParallel(modelG, list(range(ngpu))) modelD = nn.DataParallel(modelD, list(range(ngpu))) DNet = nn.DataParallel(DNet, list(range(ngpu))) modelG.apply(weights_init) modelD.apply(weights_init) criterion_perceptual = vgg_loss(device) criterion_l1 = nn.L1Loss(size_average=None, reduce=None, reduction='sum') criterion_bce = nn.BCELoss() criterion_l2 = sum_squared_error() criterion_ssim = SSIM() if torch.cuda.is_available(): modelG.to(device) modelD.to(device) DNet.to(device) optimizerG = optim.Adam(modelG.parameters(), lr=lr, betas=(0.5, 0.999), weight_decay=1e-5) optimizerD = optim.Adam(modelD.parameters(), lr=lr, betas=(0.5, 0.999), weight_decay=1e-5) scheduler = MultiStepLR(optimizerG, milestones=[30, 60, 90], gamma=0.2) # learning rates if sigma == 0: sigma_list = [5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70] else: sigma_list = [sigma] for epoch in range(n_epoch): for sig in sigma_list: x = dg.datagenerator(data_dir=data_dir).astype('float32') / 255.0 print(x.shape) x = torch.from_numpy(x.transpose((0, 3, 1, 2))) dataset = DenoisingDataset(x, sigma=sig) loader = DataLoader(dataset=dataset, num_workers=4, drop_last=True, batch_size=batch_size, shuffle=True) epoch_loss_g = 0 start_time = time.time() n_count = 0 for cnt, batch_yx in enumerate(loader): if torch.cuda.is_available(): batch_original, batch_noise = batch_yx[1].to( device), batch_yx[0].to(device) ''' modelD.zero_grad() b_size = batch_original.size(0) label = torch.full((b_size,), 1, device=device) output = modelD(batch_original).view(-1) errD_real = criterion_bce(output, label) errD_real.backward(retain_graph=True) ''' residual = DNet(batch_noise) fake, structure, denoised = modelG(batch_noise, residual) ''' label.fill_(0) output = modelD(fake.detach()).view(-1) errD_fake = criterion_bce(output, label) errD_fake.backward(retain_graph=True) d_loss = errD_real + errD_fake optimizerD.step() ''' modelG.zero_grad() ''' label = torch.full((b_size,), 1, device=device) output = modelD(fake).view(-1) gan_loss = criterion_bce(output, label) ''' s_loss = criterion_l2(structure, batch_original - denoised) s_loss.backward(retain_graph=True) l1_loss = criterion_l1(fake, batch_original) perceptual_loss = criterion_perceptual(fake, batch_original, 0) #ssim_out = 1-criterion_ssim(fake, batch_original) #l2_loss = criterion_l2(fake, batch_original) g_loss = l1_loss + 2e-2 * perceptual_loss #+1e-2*gan_loss g_loss.backward(retain_graph=True) epoch_loss_g += g_loss.item() optimizerG.step() if cnt % 100 == 0: line = '%4d %4d / %4d g_loss = %2.4f\t(snet_l2_loss = %2.4f / l1_loss=%2.4f / perceptual_loss=%2.4f)' % ( epoch + 1, cnt, x.size(0) // batch_size, g_loss.item() / batch_size, s_loss.item() / batch_size, l1_loss.item() / batch_size, perceptual_loss.item() / batch_size) print(line) f.write(line) f.write('\n') n_count += 1 elapsed_time = time.time() - start_time line = 'epoch = %4d, sigma = %4d, loss = %4.4f , time = %4.2f s' % ( epoch + 1, sig, epoch_loss_g / (n_count * batch_size), elapsed_time) print(line) f.write(line) f.write('\n') if (epoch + 1) % 20 == 0: torch.save( modelG, os.path.join( save_dir, save_name.replace('.pth', '_epoch%03d.pth') % (epoch + 1))) torch.save(modelG, os.path.join(save_dir, save_name)) f.close()
def pretrain_DNet(batch_size=128, n_epoch=150, sigma=25, lr=1e-3, depth=17, device="cuda:0", data_dir='./data/Train400', model_dir='models'): device = torch.device(device) model_dir = os.path.join(model_dir, "DNet_s" + str(sigma) + "d" + str(depth)) if not os.path.exists(model_dir): os.mkdir(model_dir) from datetime import date save_name = "DNet_s" + str(sigma) + "d" + str(depth) + "_" + "".join( str(date.today()).split('-')[1:]) + ".pth" print('\n') print('--\t This model is pre-trained DNet saved as ', save_name) print('--\t epoch %4d batch_size %4d sigma %4d depth %4d' % (n_epoch, batch_size, sigma, depth)) print('\n') model = DNet(depth=depth) model.train() print(model) print("\n") criterion = sum_squared_error() if torch.cuda.is_available(): model.to(device) optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5) scheduler = MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.2) # learning rates if sigma == 0: sigma_list = [5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70] else: sigma_list = [sigma] for epoch in range(n_epoch): for sig in sigma_list: x = dg.datagenerator(data_dir=data_dir).astype('float32') / 255.0 x = torch.from_numpy(x.transpose((0, 3, 1, 2))) dataset = None dataset = DenoisingDataset(x, sigma) loader = DataLoader(dataset=dataset, num_workers=4, drop_last=True, batch_size=batch_size, shuffle=True) epoch_loss = 0 start_time = time.time() n_count = 0 for cnt, batch_yx in enumerate(loader): optimizer.zero_grad() if torch.cuda.is_available(): batch_original, batch_noise = batch_yx[1].to( device), batch_yx[0].to(device) r = model(batch_noise) loss = criterion(batch_noise - r, batch_original) epoch_loss += loss.item() loss.backward() optimizer.step() if cnt % 100 == 0: print('%4d %4d / %4d loss = %2.4f' % (epoch + 1, cnt, x.size(0) // batch_size, loss.item() / batch_size)) n_count += 1 elapsed_time = time.time() - start_time print('epoch = %4d , sigma = %4d, loss = %4.4f , time = %4.2f s' % (epoch + 1, sig, epoch_loss / n_count, elapsed_time)) if (epoch + 1) % 25 == 0: torch.save( model, os.path.join( model_dir, save_name.replace('.pth', '_epoch%03d.pth') % (epoch + 1))) torch.save(model, os.path.join(model_dir, save_name)) torch.save(model, os.path.join(model_dir, save_name))
def pretrain_SNet(batch_size=128, n_epoch=100, sigma=25, lr=1e-4, device="cuda:0", data_dir='./data/Train400', model_dir='models/SNet', model_name=None, model=0): device = torch.device(device) if not os.path.exists(model_dir): os.mkdir(os.path.join(model_dir)) DNet = torch.load(os.path.join(model_dir, model_name[0])) if model == 0: model = SNet_jfver1() save_name = 'SNet_jfver1' elif model == 1: model = SNet_dfver1() save_name = 'SNet_dfver1' elif model == 2: model = SNet_dfver2() save_name = 'SNet_dfver2' elif model == 3: model = SNet_texture_ver1() save_name = 'SNet_texture_ver1' from datetime import date save_name = save_name + "_" + "".join(str( date.today()).split('-')[1:]) + ".pth" f = open(os.path.join(model_dir, save_name.replace(".pth", ".txt")), 'w') print('\n') print('--\t This model is pre-trained SNet saved as ', save_name) print('--\t epoch %4d batch_size %4d sigma %4d' % (n_epoch, batch_size, sigma)) print('\n') DNet.eval() model.train() print(model) print("\n") criterion = sum_squared_error() if torch.cuda.is_available(): DNet.to(device) model.to(device) optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5) scheduler = MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.2) # learning rates for epoch in range(n_epoch): x = dg.datagenerator(data_dir=data_dir).astype('float32') / 255.0 x = torch.from_numpy(x.transpose((0, 3, 1, 2))) dataset = DenoisingDataset(x, sigma) loader = DataLoader(dataset=dataset, num_workers=4, drop_last=True, batch_size=batch_size, shuffle=True) epoch_loss = 0 start_time = time.time() n_count = 0 for cnt, batch_yx in enumerate(loader): optimizer.zero_grad() if torch.cuda.is_available(): batch_original, batch_noise = batch_yx[1].to( device), batch_yx[0].to(device) r = DNet(batch_noise) d = batch_noise - r #r=1.55*(r+0.5)-0.8 s = model(r, d) #target = 1.8*(batch_original-d+0.5)-0.8 loss = criterion(s, batch_original - d) epoch_loss += loss.item() loss.backward() optimizer.step() if cnt % 100 == 0: line = '%4d %4d / %4d loss = %2.4f' % ( epoch + 1, cnt, x.size(0) // batch_size, loss.item() / batch_size) print(line) f.write(line) n_count += 1 elapsed_time = time.time() - start_time line = 'epoch = %4d , loss = %4.4f , time = %4.2f s' % ( epoch + 1, epoch_loss / n_count, elapsed_time) print(line) f.write(line) if (epoch + 1) % 1 == 0: torch.save( model, os.path.join( model_dir, save_name.replace('.pth', '_epoch%03d.pth') % (epoch + 1))) torch.save(model, os.path.join(model_dir, save_name)) f.close()
batch_y = ys[indices[i:i+batch_size]] yield batch_y, batch_x # define loss def sum_squared_error(y_true, y_pred): #return K.mean(K.square(y_pred - y_true), axis=-1) #return K.sum(K.square(y_pred - y_true), axis=-1)/2 return K.sum(K.square(y_pred - y_true))/2 if __name__ == '__main__': import pudb; pu.db # model selection model = DnCNN(depth=17,filters=64,image_channels=1,use_bnorm=True) model.summary() xs, ys = dg.datagenerator(args.train_data) # load the last model in matconvnet style initial_epoch = findLastCheckpoint(save_dir=save_dir) if initial_epoch > 0: print('resuming by loading epoch %03d'%initial_epoch) model = load_model(os.path.join(save_dir,'model_%03d.hdf5'%initial_epoch), compile=False) # compile the model model.compile(optimizer=Adam(0.001), loss=sum_squared_error) # use call back functions checkpointer = ModelCheckpoint(os.path.join(save_dir,'model_{epoch:03d}.hdf5'), verbose=1, save_weights_only=False, period=args.save_every) csv_logger = CSVLogger(os.path.join(save_dir,'log.csv'), append=True, separator=',') lr_scheduler = LearningRateScheduler(lr_schedule)
model.train() criterion = sum_squared_error() if cuda: model = model.cuda() optimizer = optim.Adam(model.parameters(), lr=args.lr) scheduler = MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.2) # learning rates for epoch in range(initial_epoch, n_epoch): for subepoch in range(0, 10): logging.info('epoch' + str(epoch)) scheduler.step(epoch) # step to the learning rate in this epcoh xs, xn = dg.datagenerator(data_dir=args.train_data, data_dir_noise=args.train_data_noise, batch_size=batch_size) listr = list(range(0, xs.shape[0])) random.shuffle(listr) xs = xs[listr, :, :, :] xn = xn[listr, :, :, :] xs = xs.astype('float32') / 255.0 xn = xn.astype('float32') / 255.0 xs = torch.from_numpy(xs.transpose( (0, 3, 1, 2))) # tensor of the clean patches, NXCXHXW xn = torch.from_numpy(xn.transpose( (0, 3, 1, 2))) # tensor of the clean patches, NXCXHXW DDataset = DenoisingDataset(xs, xn) DLoader = DataLoader(dataset=DDataset, num_workers=8, drop_last=True,
def train_model(config): # Define hyper-parameters. depth = int(config["DnCNN"]["depth"]) n_channels = int(config["DnCNN"]["n_channels"]) img_channel = int(config["DnCNN"]["img_channel"]) kernel_size = int(config["DnCNN"]["kernel_size"]) use_bnorm = config.getboolean("DnCNN", "use_bnorm") epochs = int(config["DnCNN"]["epoch"]) batch_size = int(config["DnCNN"]["batch_size"]) train_data_dir = config["DnCNN"]["train_data_dir"] test_data_dir = config["DnCNN"]["test_data_dir"] eta_min = float(config["DnCNN"]["eta_min"]) eta_max = float(config["DnCNN"]["eta_max"]) dose = float(config["DnCNN"]["dose"]) model_save_dir = config["DnCNN"]["model_save_dir"] # Save logs to txt file. log_dir = config["DnCNN"]["log_dir"] log_dir = Path(log_dir) / "dose{}".format(str(int(dose * 100))) log_file = log_dir / "train_result.txt" # Define device. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Initiate a DnCNN instance. # Load the model to device and set the model to training. model = DnCNN(depth=depth, n_channels=n_channels, img_channel=img_channel, use_bnorm=use_bnorm, kernel_size=kernel_size) model = model.to(device) model.train() # Define loss criterion and optimizer optimizer = optim.Adam(model.parameters(), lr=1e-3) scheduler = MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.2) criterion = LossFunc(reduction="mean") # Get a validation test set and corrupt with noise for validation performance. # For every epoch, use this pre-determined noisy images. test_file_list = glob.glob(test_data_dir + "/*.png") xs_test = [] # Can't directly convert the xs_test from list to ndarray because some images are 512*512 # while the rest are 256*256. for i in range(len(test_file_list)): img = cv2.imread(test_file_list[i], 0) img = np.array(img, dtype="float32") / 255.0 img = np.expand_dims(img, axis=0) img_noisy, _ = nm(img, eta_min, eta_max, dose, t=100) xs_test.append((img_noisy, img)) # Train the model. loss_store = [] epoch_loss_store = [] psnr_store = [] ssim_store = [] psnr_tr_store = [] ssim_tr_store = [] loss_mse = torch.nn.MSELoss() dtype = torch.cuda.FloatTensor # load vgg network vgg = Vgg16().type(dtype) for epoch in range(epochs): # For each epoch, generate clean augmented patches from the training directory. # Convert the data from uint8 to float32 then scale them to make it in [0, 1]. # Then make the patches to be of shape [N, C, H, W], # where N is the batch size, C is the number of color channels. # H and W are height and width of image patches. xs = dg.datagenerator(data_dir=train_data_dir) xs = xs.astype("float32") / 255.0 xs = torch.from_numpy(xs.transpose((0, 3, 1, 2))) train_set = dg.DenoisingDatatset(xs, eta_min, eta_max, dose) train_loader = DataLoader(dataset=train_set, num_workers=4, drop_last=True, batch_size=batch_size, shuffle=True) # TODO: if drop_last=True, the dropping in the # TODO: data_generator is not necessary? # train_loader_test = next(iter(train_loader)) t_start = timer() epoch_loss = 0 for idx, data in enumerate(train_loader): inputs, labels = data inputs, labels = inputs.to(device), labels.to(device) img_batch_read = len(inputs) optimizer.zero_grad() outputs = model(inputs) # We can use labels for both style and content image # style image # style_transform = transforms.Compose([ # normalize_tensor_transform() # normalize with ImageNet values # ]) # labels_t = style_transform(labels) labels_t = labels.repeat(1, 3, 1, 1) outputs_t = outputs.repeat(1, 3, 1, 1) y_c_features = vgg(labels_t) style_gram = [gram(fmap) for fmap in y_c_features] y_hat_features = vgg(outputs_t) y_hat_gram = [gram(fmap) for fmap in y_hat_features] # calculate style loss style_loss = 0.0 for j in range(4): style_loss += loss_mse(y_hat_gram[j], style_gram[j][:img_batch_read]) style_loss = STYLE_WEIGHT*style_loss aggregate_style_loss = style_loss # calculate content loss (h_relu_2_2) recon = y_c_features[1] recon_hat = y_hat_features[1] content_loss = CONTENT_WEIGHT*loss_mse(recon_hat, recon) aggregate_content_loss = content_loss loss = aggregate_content_loss + aggregate_style_loss # loss = criterion(outputs, labels) loss_store.append(loss.item()) epoch_loss += loss.item() loss.backward() optimizer.step() if idx % 100 == 0: print("Epoch [{} / {}], step [{} / {}], loss = {:.5f}, lr = {:.6f}, elapsed time = {:.2f}s".format( epoch + 1, epochs, idx, len(train_loader), loss.item(), *scheduler.get_last_lr(), timer()-t_start)) epoch_loss_store.append(epoch_loss / len(train_loader)) # At each epoch validate the result. model = model.eval() # # Firstly validate on training sets. This takes a long time so I commented. # tr_psnr = [] # tr_ssim = [] # # t_start = timer() # with torch.no_grad(): # for idx, train_data in enumerate(train_loader): # inputs, labels = train_data # # print(inputs.shape) # # inputs = np.expand_dims(inputs, axis=0) # # inputs = torch.from_numpy(inputs).to(device) # inputs = inputs.to(device) # labels = labels.squeeze().numpy() # # outputs = model(inputs) # outputs = outputs.squeeze().cpu().detach().numpy() # # tr_psnr.append(peak_signal_noise_ratio(labels, outputs)) # tr_ssim.append(structural_similarity(outputs, labels)) # psnr_tr_store.append(sum(tr_psnr) / len(tr_psnr)) # ssim_tr_store.append(sum(tr_ssim) / len(tr_ssim)) # # print("Elapsed time = {}".format(timer() - t_start)) # # print("Validation on train set: epoch [{} / {}], aver PSNR = {:.2f}, aver SSIM = {:.4f}".format( # epoch + 1, epochs, psnr_tr_store[-1], ssim_tr_store[-1])) # Validate on test set val_psnr = [] val_ssim = [] with torch.no_grad(): for idx, test_data in enumerate(xs_test): inputs, labels = test_data inputs = np.expand_dims(inputs, axis=0) inputs = torch.from_numpy(inputs).to(device) labels = labels.squeeze() outputs = model(inputs) outputs = outputs.squeeze().cpu().detach().numpy() val_psnr.append(peak_signal_noise_ratio(labels, outputs)) val_ssim.append(structural_similarity(outputs, labels)) psnr_store.append(sum(val_psnr) / len(val_psnr)) ssim_store.append(sum(val_ssim) / len(val_ssim)) print("Validation on test set: epoch [{} / {}], aver PSNR = {:.2f}, aver SSIM = {:.4f}".format( epoch + 1, epochs, psnr_store[-1], ssim_store[-1])) # Set model to train mode again. model = model.train() scheduler.step() # Save model save_model(model, model_save_dir, epoch, dose * 100) # Save the loss and validation PSNR, SSIM. if not log_dir.exists(): Path.mkdir(log_dir) with open(log_file, "a+") as fh: # fh.write("{} Epoch [{} / {}], loss = {:.6f}, train PSNR = {:.2f}dB, train SSIM = {:.4f}, " # "validation PSNR = {:.2f}dB, validation SSIM = {:.4f}".format( # datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"), # epoch + 1, epochs, epoch_loss_store[-1], # psnr_tr_store[-1], ssim_tr_store[-1], # psnr_store[-1], ssim_store[-1])) fh.write("{} Epoch [{} / {}], loss = {:.6f}, " "validation PSNR = {:.2f}dB, validation SSIM = {:.4f}\n".format( datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"), epoch + 1, epochs, epoch_loss_store[-1], psnr_store[-1], ssim_store[-1])) # np.savetxt(log_file, np.hstack((epoch + 1, epoch_loss_store[-1], psnr_store[-1], ssim_store[-1])), fmt="%.6f", delimiter=", ") fig, ax = plt.subplots() ax.plot(loss_store[-len(train_loader):]) ax.set_title("Last 1862 losses") ax.set_xlabel("iteration") fig.show()
def train(batch_size=128, n_epoch=150, sigma=25, lr=1e-3, lr2=1e-5, depth=17, device="cuda:0", data_dir='./data/Train400', model_dir='models'): device = torch.device(device) from datetime import date model_dir = os.path.join( model_dir, "result_" + "".join(str(date.today()).split('-')[1:])) if not os.path.exists(model_dir): os.mkdir(model_dir) save_name = "net.pth" save_name2 = "net2.pth" print('\n') print('--\t This model is pre-trained DNet saved as ', save_name) print('--\t epoch %4d batch_size %4d sigma %4d depth %4d' % (n_epoch, batch_size, sigma, depth)) print('\n') model = DnCNN(depth=depth) model2 = DnCNN(depth=depth) model.train() model2.train() criterion = sum_squared_error() criterion_l1 = nn.L1Loss(size_average=None, reduce=None, reduction='sum') if torch.cuda.is_available(): model.to(device) model2.to(device) optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5) scheduler = MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.2) # learning rates optimizer2 = optim.Adam(model2.parameters(), lr=lr2, weight_decay=1e-5) scheduler2 = MultiStepLR(optimizer2, milestones=[30, 60, 90], gamma=0.2) # learning rates for epoch in range(n_epoch): x = dg.datagenerator(data_dir=data_dir).astype('float32') / 255.0 x = torch.from_numpy(x.transpose((0, 3, 1, 2))) dataset = None dataset = DenoisingDataset(x, sigma) loader = DataLoader(dataset=dataset, num_workers=4, drop_last=True, batch_size=batch_size, shuffle=True) epoch_loss = 0 epoch_loss_first = 0 start_time = time.time() n_count = 0 for cnt, batch_yx in enumerate(loader): optimizer.zero_grad() if torch.cuda.is_available(): batch_original, batch_noise = batch_yx[1].to( device), batch_yx[0].to(device) residual = model(batch_noise) loss_first = criterion(batch_noise - residual, batch_original) loss_first.backward(retain_graph=True) optimizer.step() structure_residual = model2(residual) target = batch_original - (batch_noise - residual) structure = residual - structure_residual loss = criterion_l1(structure, target) loss.backward() optimizer2.step() epoch_loss_first += loss_first.item() epoch_loss += loss.item() if cnt % 100 == 0: print( '%4d %4d / %4d 1_loss = %2.4f loss = %2.4f' % (epoch + 1, cnt, x.size(0) // batch_size, loss_first.item() / batch_size, loss.item() / batch_size)) n_count += 1 elapsed_time = time.time() - start_time print( 'epoch = %4d , sigma = %4d, 1_loss = %4.4f, loss = %4.4f , time = %4.2f s' % (epoch + 1, sigma, epoch_loss_first / n_count, epoch_loss / n_count, elapsed_time)) if (epoch + 1) % 25 == 0: torch.save( model, os.path.join( model_dir, save_name.replace('.pth', '_epoch%03d.pth') % (epoch + 1))) torch.save( model2, os.path.join( model_dir, save_name2.replace('.pth', '_epoch%03d.pth') % (epoch + 1))) torch.save(model, os.path.join(model_dir, save_name)) torch.save(model2, os.path.join(model_dir, save_name2))
model = torch.load( os.path.join(save_dir, 'model_%03d.pth' % initial_epoch)) model.train() # criterion = nn.MSELoss(reduction = 'sum') # PyTorch 0.4.1 criterion = sum_squard_error() if use_cuda: model = model.cuda() criterion = criterion.cuda() optimizer = optim.SGD(model.parameters(), lr=args.lr) scheduler = Polyscheduler(optimizer, EPOCH=n_epoch, gamma=0.8) # learning rates # scheduler = LrScheduler.ReduceLROnPlateau(optimizer, 'min') # learning rates writer = SummaryWriter(comment='train') DATA = dg.datagenerator(data_dir=filename) # data = torch.rand(20, 1, 11, 11) DATA = DATA.astype('float32') / 255.0 DATA = torch.from_numpy(DATA.transpose((0, 3, 1, 2))) LABEL = dg.datagenerator(data_dir=filename1) # LABEL = torch.rand(20, 1, 11, 11) LABEL = LABEL.astype('float32') / 255.0 LABEL = torch.from_numpy(LABEL.transpose((0, 3, 1, 2))) torch_dataset = DATASET(DATA, LABEL) for epoch in range(initial_epoch, n_epoch): scheduler.step(epoch) # step to the learning rate in this epcoh LR = torch.FloatTensor(scheduler.get_lr()) DLoader = DataLoader(dataset=torch_dataset, num_workers=4,