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))
random.seed(seed)
torch.backends.cudnn.deterministic = True
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,