def train(args):
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
mean_value = 0.5
# root = 'data/train'
train_dataset = Patches(root=args.path, phase='train')
trainloader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batchsize)
model = EDSR(num_layers=args.layers, feature_size=args.featuresize)
model.to(device).train()
model = nn.DataParallel(model, device_ids=range(
torch.cuda.device_count())).to(device)
model.train()
optimizer = optim.Adam(model.parameters(), lr=0.0001)
for epoch in range(1, args.epochs + 1):
pbar = tqdm(trainloader)
for data in pbar:
data50, data100 = data
data50, data100 = data50.to(device), data100.to(device)
optimizer.zero_grad()
output, _ = model(data50)
loss = F.l1_loss(output, data100 - mean_value)
loss.backward()
optimizer.step()
pbar.set_description("epoch: %d train_loss: %.4f" %
(epoch, loss.item()))
if epoch % 50 == 0:
torch.save(model.state_dict(),
args.savedir + '/edsr_step_{}.pth'.format(epoch))
def main(args):
cfg = cfg_dict[args.cfg_name]
writer = SummaryWriter(os.path.join("runs", args.cfg_name))
train_loader = get_data_loader(cfg, cfg["train_dir"])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = EDSR(cfg).to(device)
criterion = torch.nn.L1Loss()
optimizer = torch.optim.Adam(model.parameters(), lr=cfg["init_lr"],
betas=(0.9, 0.999), eps=1e-8)
global_batches = 0
if args.train:
for epoch in range(cfg["n_epoch"]):
model.train()
running_loss = 0.0
for i, batch in enumerate(train_loader):
lr, hr = batch[0].to(device), batch[1].to(device)
optimizer.zero_grad()
sr = model(lr)
loss = model.loss(sr, hr)
# loss = criterion(model(lr), hr)
running_loss += loss.item()
loss.backward()
optimizer.step()
global_batches += 1
if global_batches % cfg["lr_decay_every"] == 0:
for param_group in optimizer.param_groups:
print(f"decay lr to {param_group['lr'] / 10}")
param_group["lr"] /= 10
if epoch % args.log_every == 0:
model.eval()
with torch.no_grad():
batch_samples = {"lr": batch[0], "hr": batch[1],
"sr": sr.cpu()}
writer.add_scalar("training-loss",
running_loss / len(train_loader),
global_step=global_batches)
writer.add_scalar("PSNR", compute_psnr(batch_samples),
global_step=global_batches)
samples = {k: v[:3] for k, v in batch_samples.items()}
fig = visualize_samples(samples, f"epoch-{epoch}")
writer.add_figure("sample-visualization", fig,
global_step=global_batches)
if epoch % args.save_every == 0:
state = {"net": model.state_dict(),
"optim": optimizer.state_dict()}
checkpoint_dir = args.checkpoint_dir
if not os.path.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
path = os.path.join(checkpoint_dir, args.cfg_name)
torch.save(state, path)
# eval
if args.eval:
assert args.model_path and args.lr_img_path
print(f"evaluating {args.lr_img_path}")
state = torch.load(args.model_path, map_location=device)
model.load_state_dict(state["net"])
optimizer.load_state_dict(state["optim"])
with torch.no_grad():
lr = img2tensor(args.lr_img_path)
sr = model(lr.clone().to(device)).cpu()
samples = {"lr": lr, "sr": sr}
if args.hr_img_path:
samples["hr"] = img2tensor(args.hr_img_path)
print(f"PSNR: {compute_psnr(samples)}")
directory = os.path.dirname(args.lr_img_path)
name = f"eval-{args.cfg_name}-{args.lr_img_path.split('/')[-1]}"
visualize_samples(samples, name, save=True,
directory=directory, size=6)
def loss_fn(sr, gt):
loss = nn.MSELoss(reduction='sum')
output = loss(sr, gt)
return output
if __name__ == '__main__':
set_seed(2019) # Set seed to produce the same training results
model = EDSR(upscale=scale)
model = nn.DataParallel(model, device_ids=[0])
model = model.to(device)
model.train()
optimizer = torch.optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=weight_decay)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=lr_step,
gamma=0.1)
writer = SummaryWriter(os.path.join(checkpoint_dir, 'tensorboard_log'),
flush_secs=10)
step = 0
dataset = DIV2K_Dataset(dataset_dir, patch_size, scale, crop_num_per_image)
dl = DataLoader(dataset,
batch_size=batch_size,
shuffle=True,
num_workers=0,
drop_last=True,