예제 #1
0
# 读取数据集
train_dataset = ImagePairDataset_y(INPUTS_FOLDER_TRAIN, LABELS_FOLDER_TRAIN)
test_dataset = ImagePairDataset_y(INPUTS_FOLDER_TEST, LABELS_FOLDER_TEST)
train_iter = DataLoader(train_dataset, batch_size, shuffle=True)
test_iter = DataLoader(test_dataset, 1, shuffle=True)
logging.info('Datasets loaded!')

# 训练
train(train_iter, test_iter, net, loss, my_optim, num_epochs, scheduler)

# 测试
logging.info('Final test loss %.4f' % eval(test_iter, net, loss, 0))

# 保存网络
torch.save(net.state_dict(), SAVE_PATH)
logging.info('Network saved: ' + SAVE_PATH)

# 读取网络
# net.load_state_dict(torch.load(LOAD_PATH))
# logging.info('Network loaded: '+LOAD_PATH)

# 应用完整图片并写入
IMG_NAME = r'img_0'
IMG_OUT = OUTPUTS_FOLDER+r'/'+IMG_NAME+r'_' + \
    computer_name+r'_'+format_time+r'.png'
output = apply_net(LR_FOLDER + r'/' + IMG_NAME + r'.png', IMG_OUT, net, ratio)

# 对测试图片进行评估
img_metric(IMG_OUT, HR_FOLDER + r'/' + IMG_NAME + r'.png')
예제 #2
0
# batch_size = 64
# my_optim=torch.optim.SGD(net.parameters(),lr=lr,momentum=0.9,weight_decay=0.0001)
# # 自适应学习率
# scheduler = torch.optim.lr_scheduler.StepLR(my_optim,step_size=20,gamma = 0.1)
# loss = torch.nn.MSELoss()

# # 读取数据集
# train_dataset=ImagePairDataset_y(Inputs_folder_train,Labels_folder_train)
# test_dataset=ImagePairDataset_y(Inputs_folder_test,Labels_folder_test)
# train_iter = DataLoader(train_dataset, batch_size, shuffle=True)
# test_iter = DataLoader(test_dataset, 1, shuffle=True)
# print('Datasets loaded!')

# # 训练
# train(train_iter, test_iter, net, loss, my_optim, num_epochs,scheduler)

# # 测试
# print('Full test loss %.4f'%eval(test_iter,net,loss,0))

# # 保存网络
# torch.save(net.state_dict(), MODEL_PATH)

# # 读取网络
# net.load_state_dict(torch.load(MODEL_PATH))

# # 应用完整图片并写入
# output=apply_net(LR_folder+r'/img_1.png', Outputs_folder+r'/img_1_out.png',net)

# 对测试图片进行评估
img_metric(Outputs_folder + r'/img_1_out.png', HR_folder + r'/img_1.png')