def train(training_data, dev_data, args):
training_gen = data.DataLoader(training_data, batch_size=2)
dev_gen = data.DataLoader(dev_data, batch_size=2)
device = torch.device('cuda' if cuda.is_available() else 'cpu')
print('Initializing model')
model = SRCNN()
loss = RMSE()
if cuda.device_count() > 1:
print('Using %d CUDA devices' % cuda.device_count())
model = nn.DataParallel(
model, device_ids=[i for i in range(cuda.device_count())])
model.to(device)
loss.to(device)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
def _train(data, opt=True):
total = 0
for y, x in data:
y, x = y.to(device), x.to(device)
pred_y = model(x)
l = loss(pred_y, y)
total += l.item()
if opt:
optimizer.zero_grad()
l.backward()
optimizer.step()
cuda.synchronize()
return total
print('Training')
for ep in range(args.ep):
train_loss = _train(training_gen)
dev_loss = _train(dev_gen, opt=False)
print_flush('Epoch %d: Train %.4f Dev %.4f' %
(ep, train_loss, dev_loss))
if ep % 50 == 0:
save_model(model, args.o)
return model
help="use CUDA to speed up computation")
opt = parser.parse_args()
# check cuda
if opt.cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
# load LR image
open_image = lambda x: TF.to_tensor(Image.open(x).convert("RGB"))
image = open_image(opt.filename).unsqueeze(0)
# load model
device = torch.device(
"cuda" if opt.cuda and torch.cuda.is_available() else "cpu")
map_location = "cuda:0" if opt.cuda else device
checkpoint = load_checkpoint(opt.checkpoint, map_location)
# init model
model = SRCNN()
model.load_state_dict(checkpoint["model_state_dict"])
model.to(device)
model.eval()
# process input
with torch.no_grad():
image = image.to(device)
output = model(image)
# save result
output_name = os.path.splitext(opt.filename)
output_name = "_srcnn_x3".join(output_name)
torchvision.utils.save_image(output[0], output_name)
print("Output HR image saved to '{output}'".format(output=output_name))
