'params':
models['upscale'][scale].parameters(),
'lr':
LEARNING_RATE * 0.1
})
optimizer['generator'].add_param_group({
'params':
models['extra'].parameters(),
'lr':
LEARNING_RATE * 0.1
})
optimizer['discriminator'] = optim.Adam(
models['discriminator'].parameters(), lr=LEARNING_RATE)
#set Meter to calculate the average of loss
train_loss = meter.AverageMeter()
#running
for epoch in range(EPOCH_START, EPOCH_START + EPOCH):
#select upscale factor
scale = UPSCALE_FACTOR_LIST[epoch % len(UPSCALE_FACTOR_LIST)]
datasets.scale_factor = scale
for iteration in range(ITER_PER_EPOCH):
train(models, scale, train_data_loader, criterion, optimizer,
train_loss, False)
print('{:0>3d}: train_loss: {:.8f}, scale: {}'.format(
epoch + 1, train_loss.avg, scale))
#save
if PRETRAINED == True:
batch_size=BSIZE,
sampler=SubsetRandomSampler(valid_idx),
num_workers=0)
model = model.IrisNetwork(4, 32, 3)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=LRATE)
best_loss = 1.5
history = {
'epoch': [],
'train_loss': [],
'valid_loss': [],
}
for epoch in range(NUM_EPOCH):
batch_time = meter.AverageMeter()
data_time = meter.AverageMeter()
losses = meter.AverageMeter()
end_time = time.time()
for idx, (x_train, y_train) in enumerate(train_loader):
data_time.update(time.time() - end_time)
out = model(x_train)
loss = criterion(out, y_train)
loss.backward()
optimizer.step()
losses.update(loss.item(), x_train.size(0))
batch_time.update(time.time() - end_time)
end_time = time.time()
device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
#load data
test_datas=datasets.TestData[UPSCALE_FACTOR]
#add model
models={}
models['generative']=model.FSRCNN().to(device)
models['upscale']={}
models['upscale'][UPSCALE_FACTOR]=model.Subpixel_Layer(models['generative'].output_channel,UPSCALE_FACTOR).to(device)
models['extra']=model.Extra_Layer().to(device)
interpolate=False
#set criterion
criterion=nn.MSELoss()
#set meter
PSNR=meter.AverageMeter()
SSIM=meter.AverageMeter()
for key in test_datas.keys():
print(key)
test_data=test_datas[key]
test_data_loader=torch.utils.data.DataLoader(dataset=test_data,
batch_size=1,
shuffle=False)
PSNR_array=[]
SSIM_array=[]
best=0
for epoch in range(EPOCH):
utils.load_model(models,UPSCALE_FACTOR,LOAD_PATH,epoch+1)
with torch.no_grad():
test(models,UPSCALE_FACTOR,test_data_loader,criterion,PSNR,SSIM,False)