def main():
model_dir = os.path.join(opt.outf, opt.name)
print('create checkpoint directory %s...' % model_dir)
if not os.path.exists(model_dir):
os.makedirs(model_dir)
# Load dataset
print('Loading dataset ...\n')
dataset_train = Dataset(train=True)
loader_train = DataLoader(dataset=dataset_train, num_workers=4, batch_size=opt.batchSize, shuffle=True)
print("# of training samples: %d\n" % int(len(dataset_train)))
num = len(dataset_train)
# Build model
net = RDN(64, 3)
num_params = 0
for parm in net.parameters():
num_params += parm.numel()
print(net)
print('[Network %s] Total number of parameters : %.3f M' % (opt.name, num_params / 1e6))
# Move to GPU
device_ids = [0]
model = nn.DataParallel(net, device_ids=device_ids).cuda()
#model.load_state_dict(torch.load(os.path.join('logs/', opt.name, '40_net.pth'))) # !!!
# Optimizer
optimizer = optim.Adam(model.parameters(), lr=opt.lr) #weight_decay=opt.weight_decay
# training
step = 0
for epoch in range(opt.epochs):
# set learning rate
current_lr = update_lr(opt.lr, epoch)
for param_group in optimizer.param_groups:
param_group["lr"] = current_lr
print('learning rate %f' % current_lr)
# train
model.train()
start_time = time.time()
ave_loss = 0
ave_psnr = 0
ave_ssim = 0
for i, data in enumerate(loader_train, 0):
# training step
time1 = time.time()
model.zero_grad()
optimizer.zero_grad()
noise_img = data[:, :3, :, :]
gt_img = data[:, 3:, :, :]
noise_img, gt_img = noise_img.cuda(), gt_img.cuda()
res = noise_img - gt_img
pred_res = model(noise_img)
loss1 = torch.mean(torch.abs(pred_res - gt_img))
loss2 = torch.mean(SSIM(pred_res, gt_img))
loss = loss1 #0.75*loss1 + 0.25*loss2
loss.backward()
optimizer.step()
# evaluate
#result = torch.clamp(noise_img-pred_res, 0., 1.)
result = torch.clamp(pred_res, 0., 1.)
psnr_train = batch_PSNR(result, gt_img, 1.)
ave_loss = (ave_loss*i + loss.item()) / (i+1)
ave_psnr = (ave_psnr*i + psnr_train) / (i+1)
ave_ssim = (ave_ssim*i + 1-loss2.item()*2) / (i+1)
time2 = time.time()
if i % 100 == 0:
print("[epoch %d][%d/%d] time: %.3f t_time: %.3f loss: %.4f PSNR_train: %.4f SSIM_train: %.4f" %
(epoch+1, i, len(loader_train), (time2 - time1), (time2 - start_time), ave_loss, ave_psnr, ave_ssim))
if step % 1000 == 0:
torch.save(model.state_dict(), os.path.join(model_dir, 'latest_net.pth'))
step += 1
print('Time for the epoch is %f' % (time.time() - start_time))
## the end of each epoch
# save model
save_name = '%d_net.pth' % (epoch+1)
torch.save(model.state_dict(), os.path.join(model_dir, save_name))
num_channels=3,
num_features=args.num_features,
growth_rate=args.growth_rate,
num_blocks=args.num_blocks,
num_layers=args.num_layers).to(device)
if args.weights_file is not None:
state_dict = model.state_dict()
for n, p in torch.load(args.weights_file, map_location=lambda storage, loc: storage).items():
if n in state_dict.keys():
state_dict[n].copy_(p)
else:
raise KeyError(n)
criterion = nn.L1Loss()
optimizer = optim.Adam(model.parameters(), lr=args.lr)
train_dataset = TrainDataset(args.train_file, patch_size=args.patch_size, scale=args.scale)
train_dataloader = DataLoader(dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
eval_dataset = EvalDataset(args.eval_file)
eval_dataloader = DataLoader(dataset=eval_dataset, batch_size=1)
best_weights = copy.deepcopy(model.state_dict())
best_epoch = 0
best_psnr = 0.0
for epoch in range(args.num_epochs):