def predict_model(test_fnames, x_test, test_transforms, num_classes, *, tta=5):
batch_size = 64
test_dataset = dataset.TestDataset(test_fnames,
x_test,
test_transforms,
tta=tta)
test_loader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
pin_memory=True)
net = model.MainModel(model_type='Simple', num_classes=num_classes)
net = net.model.cuda()
net.load_state_dict(torch.load('weight_best.pt'))
net.cuda()
net.eval()
all_outputs, all_fnames = [], []
for images, fnames in test_loader:
preds = torch.sigmoid(net(images.cuda()).detach())
all_outputs.append(preds.cpu().numpy())
all_fnames.extend(fnames)
test_preds = pd.DataFrame(data=np.concatenate(all_outputs),
index=all_fnames,
columns=map(str, range(num_classes)))
test_preds = test_preds.groupby(level=0).mean()
return test_preds
def main():
state = torch.load(MODEL)
net = train.Net()
net.load_state_dict(state)
data = dataset.TestDataset(start=2639, size=267)
# data = dataset.TestDataset(start=2907, size=166)
# data = dataset.TestDataset(start=3073, size=166)
# data = dataset.TestDataset(start=2639, size=600)
loader = DataLoader(data, batch_size=1)
calc_score(loader, net, "cpu")
def main(
experiment_name,
dropout_p,
batch_size,
num_workers,
augmentation,
folds,
):
transforms = make_augmentation_transforms(augmentation, mode='test')
test_dataset = dataset.TestDataset(transform=transforms)
model = make_nasnet_model(
num_classes=config.NUM_CLASSES,
dropout_p=dropout_p,
)
test_data_loader = DataLoader(
dataset=test_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=False,
pin_memory=True,
)
test_predictions = np.zeros((folds, len(test_dataset), config.NUM_CLASSES))
for fold_num in range(folds):
checkpoint = load_checkpoint(
f'{experiment_name}_{fold_num}_{folds}_best.pth')
model.load_state_dict(checkpoint['state_dict'])
model = model.cuda().eval()
with torch.no_grad():
for batch_index, (waves, _) in enumerate(tqdm(test_data_loader)):
waves = Variable(waves).cuda()
logits = model(waves)
probs = F.softmax(logits, dim=1)
numpy_probs = probs.cpu().data.numpy()
start_index = batch_index * batch_size
end_index = start_index + numpy_probs.shape[0]
test_predictions[fold_num, start_index:end_index] = numpy_probs
save_predictions(test_predictions, f'{experiment_name}.pkl')
def main():
global best_acc
start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch
if not os.path.isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# Data
transform = get_transforms(input_size=args.image_size,
test_size=args.image_size,
backbone=None)
print('==> Preparing dataset %s' % args.trainroot)
trainset = dataset.Dataset(root=args.trainroot,
transform=transform['val_train'])
train_loader = data.DataLoader(trainset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
valset = dataset.TestDataset(root=args.valroot,
transform=transform['val_test'])
val_loader = data.DataLoader(valset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
model = make_model(args)
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = get_optimizer(model, args)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
mode='min',
factor=0.2,
patience=5,
verbose=False)
# Resume
title = 'ImageNet-' + args.arch
if args.resume:
# Load checkpoint.
print('==> Resuming from checkpoint..')
assert os.path.isfile(
args.resume), 'Error: no checkpoint directory found!'
args.checkpoint = os.path.dirname(args.resume)
checkpoint = torch.load(args.resume)
best_acc = checkpoint['best_acc']
start_epoch = checkpoint['epoch']
model.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger = Logger(os.path.join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
logger.set_names([
'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.',
'Valid Acc.'
])
if args.evaluate:
print('\nEvaluation only')
test_loss, test_acc = test(val_loader, model, criterion, start_epoch,
use_cuda)
print(' Test Loss: %.8f, Test Acc: %.2f' % (test_loss, test_acc))
return
# Train and val
for epoch in range(start_epoch, args.epochs):
print('\nEpoch: [%d | %d] LR: %f' %
(epoch + 1, args.epochs, optimizer.param_groups[0]['lr']))
train_loss, train_acc, train_5 = train(train_loader, model, criterion,
optimizer, epoch, use_cuda)
test_loss, test_acc, test_5 = test(val_loader, model, criterion, epoch,
use_cuda)
scheduler.step(test_loss)
# append logger file
logger.append(
[state['lr'], train_loss, test_loss, train_acc, test_acc])
print(
'train_loss:%f, val_loss:%f, train_acc:%f, train_5:%f, val_acc:%f, val_5:%f'
% (train_loss, test_loss, train_acc, train_5, test_acc, test_5))
# save model
is_best = test_acc > best_acc
best_acc = max(test_acc, best_acc)
if len(args.gpu_id) > 1:
save_checkpoint(
{
'fold': 0,
'epoch': epoch + 1,
'state_dict': model.module.state_dict(),
'train_acc': train_acc,
'acc': test_acc,
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
is_best,
single=True,
checkpoint=args.checkpoint)
else:
save_checkpoint(
{
'fold': 0,
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'train_acc': train_acc,
'acc': test_acc,
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
is_best,
single=True,
checkpoint=args.checkpoint)
logger.close()
logger.plot()
savefig(os.path.join(args.checkpoint, 'log.eps'))
print('Best acc:')
print(best_acc)
def main():
# Data
TRAIN = args.trainroot
VAL = args.valroot
# TRAIN = '/content/train'
# VAL = '/content/val'
transform = get_transforms(input_size=args.image_size, test_size=args.image_size, backbone=None)
print('==> Preparing dataset %s' % args.trainroot)
# trainset = datasets.ImageFolder(root=TRAIN, transform=transform['train'])
# valset = datasets.ImageFolder(root=VAL, transform=transform['val'])
trainset = dataset.Dataset(root=args.trainroot, transform=transform['train'])
valset = dataset.TestDataset(root=args.valroot, transform=transform['val'])
train_loader = DataLoader(
trainset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = DataLoader(
valset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# model initial
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = make_model(args)
# TODO:merge in function
for k,v in model.named_parameters():
# print("{}: {}".format(k,v.requires_grad))
if not k.startswith('layer4') and not k.startswith('fc'):
# print(k)
v.requires_grad = False
# sys.exit(0)
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
elif args.ngpu:
model = torch.nn.DataParallel(model).cuda()
model.to(device)
cudnn.benchmark = True
print('Total params:%.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) # 打印模型参数数量
print('Trainable params:%.2fM' % (sum(p.numel() for p in model.parameters() if p.requires_grad)/ 1000000.0))
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = get_optimizer(model, args)
# 基于标准的学习率更新
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, patience=5, verbose=False)
# Resume
epochs = args.epochs
start_epoch = args.start_epoch
title = 'log-' + args.arch
if args.resume:
# --resume checkpoint/checkpoint.pth.tar
# load checkpoint
print('Resuming from checkpoint...')
assert os.path.isfile(args.resume), 'Error: no checkpoint directory found!!'
checkpoint = torch.load(args.resume)
best_acc = checkpoint['best_acc']
start_epoch = checkpoint['epoch']
state_dict = checkpoint['state_dict']
optim = checkpoint['optimizer']
model.load_state_dict(state_dict, strict=False)
optimizer.load_state_dict(optim)
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title, resume=True)
else:
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
# logger.set_names(['Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.', 'Valid Acc.'])
logger.set_names(['LR', 'epoch', 'Train Loss', 'Valid Loss', 'Train Acc.', 'Valid Acc.',])
# Evaluation:Confusion Matrix:Precision Recall F1-score
if args.evaluate and args.resume:
print('\nEvaluate only')
test_loss, test_acc, test_acc_5, predict_all, labels_all = test_model(val_loader, model, criterion, device, test=True)
print('Test Loss:%.8f,Test top1:%.2f top5:%.2f' %(test_loss,test_acc,test_acc_5))
# 混淆矩阵
report = metrics.classification_report(labels_all,predict_all,target_names=class_list,digits=4)
confusion = metrics.confusion_matrix(labels_all,predict_all)
print('\n report ',report)
print('\n confusion',confusion)
with open(args.resume[:-3]+"txt", "w+") as f_obj:
f_obj.write(report)
# plot_Matrix(args.resume[:-3], confusion, class_list)
return
# model train and val
best_acc = 0
for epoch in range(start_epoch, epochs + 1):
print('[{}/{}] Training'.format(epoch, args.epochs))
# train
train_loss, train_acc, train_acc_5 = train_model(train_loader, model, criterion, optimizer, device)
# val
test_loss, test_acc, test_acc_5 = test_model(val_loader, model, criterion, device, test=None)
scheduler.step(test_loss)
lr_ = optimizer.param_groups[0]['lr']
# 核心参数保存logger
logger.append([lr_, int(epoch), train_loss, test_loss, train_acc, test_acc,])
print('train_loss:%f, val_loss:%f, train_acc:%f, val_acc:%f, train_acc_5:%f, val_acc_5:%f' % (train_loss, test_loss, train_acc, test_acc, train_acc_5, test_acc_5))
# 保存模型 保存最优
is_best = test_acc > best_acc
best_acc = max(test_acc, best_acc)
if not args.ngpu:
name = 'checkpoint_' + str(epoch) + '.pth.tar'
else:
name = 'ngpu_checkpoint_' + str(epoch) + '.pth.tar'
save_checkpoint({
'epoch': epoch,
'state_dict': model.state_dict(),
'train_acc': train_acc,
'test_acc': test_acc,
'best_acc': best_acc,
'optimizer': optimizer.state_dict()
}, is_best, checkpoint=args.checkpoint, filename=name)
# logger.close()
# logger.plot()
# savefig(os.path.join(args.checkpoint, 'log.eps'))
print('Best acc:')
print(best_acc)
def main():
parser = argparse.ArgumentParser()
add_arguments(parser)
args = parser.parse_args()
run_dir = Path(args.run_dir)
assert torch.cuda.is_available(), 'CUDA is not available'
train_valid_transform_1 = transforms.Compose([
RandomCrop(args.input_size),
RandomRotation(), # x4
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
# train_valid_transform_2 = transforms.Compose([
# CenterCrop(args.input_size),
# RandomRotation(), # x4
# transforms.ToTensor(),
# transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
# ])
def init_loaders(transform):
fix_path = utils.fix_jpg_tif
# do_manip gives x8 samples
train_dataset = D.ConcatDataset([
dataset.CSVDataset(dataset.TRAINVAL_SET,
args,
transform=transform,
do_manip=True,
repeats=1,
fix_path=fix_path),
dataset.CSVDataset(dataset.FLICKR_TRAIN_SET,
args,
transform=transform,
do_manip=True,
repeats=1,
fix_path=fix_path)
])
valid_dataset = D.ConcatDataset([
dataset.CSVDataset(dataset.FLICKR_VALID_SET,
args,
transform=transform,
do_manip=True,
repeats=1,
fix_path=fix_path),
dataset.CSVDataset(dataset.REVIEWS_SET,
args,
transform=transform,
do_manip=True,
repeats=1,
fix_path=fix_path)
])
train_loader = D.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
valid_loader = D.DataLoader(valid_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
return train_loader, valid_loader
train_loader_1, valid_loader_1 = init_loaders(train_valid_transform_1)
# train_loader_2, valid_loader_2 = init_loaders(train_valid_transform_2)
test_transform = transforms.Compose([
CenterCrop(args.input_size),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
test_dataset = dataset.TestDataset(transform=test_transform)
test_loader = D.DataLoader(test_dataset,
batch_size=args.batch_size,
num_workers=args.workers)
model = getattr(models, args.model)(num_classes=dataset.NUM_CLASSES,
pretrained=True)
model = N.DataParallel(model).cuda()
loss = N.CrossEntropyLoss()
if 'train' == args.mode:
if run_dir.exists() and args.clean:
shutil.rmtree(str(run_dir))
run_dir.mkdir(exist_ok=True)
run_dir.joinpath('params.json').write_text(
json.dumps(vars(args), indent=True, sort_keys=True))
subprocess.check_call("git diff $(find . -name '*.py') > {}".format(
run_dir / 'patch'),
shell=True)
def init_optimizer(lr, lr_changes):
return O.Adam(model.parameters(), lr=lr)
# return O.SGD([{'params': model.feature_parameters(), 'lr': lr},
# {'params': model.classifier_parameters(), 'lr': 1e-6}], momentum=0.9)
# if lr_changes == 0:
# return O.SGD([{'params': model.module.feature_parameters(), 'lr': lr},
# {'params': model.module.classifier_parameters(), 'lr': 1e-6}],
# momentum=0.9)
# else:
# return O.SGD([{'params': model.module.feature_parameters(), 'lr': lr},
# {'params': model.module.classifier_parameters(), 'lr': 1e-4}],
# momentum=0.9)
train_kwargs = {
'args': args,
'model': model,
'criterion': loss,
}
# Train on random crops on full image
train(
init_optimizer=init_optimizer,
lr=args.lr,
train_loader=train_loader_1,
valid_loader=valid_loader_1,
# n_epochs=20,
**train_kwargs)
# Train on central crops
# train(
# init_optimizer=init_optimizer,
# lr=args.lr,
# train_loader=train_loader_2,
# valid_loader=valid_loader_2,
# **train_kwargs)
elif args.mode in ['valid', 'predict_valid', 'predict_test']:
if 'best' == args.checkpoint:
ckpt = (run_dir / 'best-model.pt')
else:
ckpt = (run_dir / 'model.pt')
state = torch.load(str(ckpt))
model.load_state_dict(state['model'])
print('Loaded {}'.format(ckpt))
if 'valid' == args.mode:
validation(tqdm.tqdm(valid_loader_1, desc='Validation'), model,
loss)
elif 'predict_valid' == args.mode:
preds, targets, manips = predict_valid(valid_loader_2, model)
save_valid_predictions(preds, targets, manips, args)
elif 'predict_test' == args.mode:
preds, paths = predict_test(test_loader, model, args.tta)
save_test_predictions(preds, paths, args)
else:
raise ValueError('Unknown mode {}'.format(args.mode))
# 3개
model_age = model_functions.load_models(
'effi-b7', "/opt/ml/model_location/0408/efficientb7_test/age/1",
"best_model.pth")
model_gender = model_functions.load_models(
'effi-b6', "/opt/ml/model_location/0408/efficientb7_test/gender/1",
"best_model.pth")
model_mask = model_functions.load_models(
'effi-b4', "/opt/ml/model_location/0408/efficientb7_test/mask/1",
"best_model.pth")
# Test Dataset 클래스 객체를 생성하고 DataLoader를 만듭니다.
test_csv = dataset.test_csv
image_paths = dataset.image_paths
transformation = dataset.make_transform(transform_type=('test'))
test_dataset = dataset.TestDataset(image_paths, transformation['test'])
loader = DataLoader(test_dataset, shuffle=False)
# 모델이 테스트 데이터셋을 예측하고 결과를 저장합니다.
def make_prediction(model_su, model_list, csv_file, data_loader):
all_predictions = []
# 단일 모델로 클래스 18개 분류할 때
if model_su == 1:
model = model_list[0]
model.cuda().float()
model.eval()
for images in tqdm(data_loader):
with torch.no_grad():
print(opt)
if opt.cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
print('===> Loading datasets')
root_path = '/'
train_set_path = os.path.join(root_path, 'train_set')
prior_set_path = os.path.join(root_path, 'prior_set')
test_set_path = os.path.join(root_path, 'test_set')
train_set = dataset.TrainDataset(train_set_path,
filenumber=11,
sample_number=1000)
train_unpair_set = dataset.PriorDataset(prior_set_path, sample_number=1000)
test_set = dataset.TestDataset(test_set_path)
training_iterator = DataLoader(dataset=train_set,
batch_size=opt.batch_size,
num_workers=opt.threads,
shuffle=True)
train_unpair_iterator = DataLoader(dataset=train_unpair_set,
batch_size=opt.batch_size,
num_workers=opt.threads,
shuffle=True)
test_iterator = DataLoader(dataset=test_set,
batch_size=opt.batch_size,
num_workers=opt.threads,
shuffle=True)
device = torch.device("cuda:0" if opt.cuda else "cpu")
def MyDNN(opt):
# ----------------------------------------
# Network training parameters
# ----------------------------------------
# cudnn benchmark
cudnn.benchmark = opt.cudnn_benchmark
# configurations
save_folder = os.path.join(opt.save_path, opt.task)
sample_folder = os.path.join(opt.sample_path, opt.task)
if not os.path.exists(save_folder):
os.makedirs(save_folder)
if not os.path.exists(sample_folder):
os.makedirs(sample_folder)
# Loss functions
criterion_L1 = torch.nn.L1Loss().cuda()
criterion_L2 = torch.nn.MSELoss().cuda()
mse_loss = nn.MSELoss().cuda()
ms_ssim_module = MS_SSIM(data_range=2,
size_average=True,
channel=3,
nonnegative_ssim=True)
# Pretrained VGG
# vgg = MINCFeatureExtractor(opt).cuda()
# Initialize Generator
generator = utils.create_MyDNN(opt)
use_checkpoint = False
if use_checkpoint:
checkpoint_path = './MyDNN1_denoise_epoch175_bs1'
# Load a pre-trained network
pretrained_net = torch.load(checkpoint_path + '.pth')
load_dict(generator, pretrained_net)
print('Generator is loaded!')
# To device
if opt.multi_gpu:
generator = nn.DataParallel(generator)
generator = generator.cuda()
else:
generator = generator.cuda()
# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(),
lr=opt.lr_g,
betas=(opt.b1, opt.b2),
weight_decay=opt.weight_decay)
# Learning rate decrease
def adjust_learning_rate(opt, epoch, iteration, optimizer):
#Set the learning rate to the initial LR decayed by "lr_decrease_factor" every "lr_decrease_epoch" epochs
if opt.lr_decrease_mode == 'epoch':
lr = opt.lr_g * (opt.lr_decrease_factor
**(epoch // opt.lr_decrease_epoch))
if epoch < 200:
lr = 0.0001
if epoch >= 200:
lr = 0.00005
if epoch >= 300:
lr = 0.00001
for param_group in optimizer.param_groups:
param_group['lr'] = lr
if opt.lr_decrease_mode == 'iter':
lr = opt.lr_g * (opt.lr_decrease_factor
**(iteration // opt.lr_decrease_iter))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
return lr
# Save the model if pre_train == True
def save_model(opt, epoch, iteration, len_dataset, generator, val_PSNR,
best_PSNR):
"""Save the model at "checkpoint_interval" and its multiple"""
if opt.save_best_model and best_PSNR == val_PSNR:
torch.save(generator,
'final_%s_epoch%d_best.pth' % (opt.task, epoch))
print('The best model is successfully saved at epoch %d' % (epoch))
if opt.multi_gpu == True:
if opt.save_mode == 'epoch':
if (epoch % opt.save_by_epoch
== 0) and (iteration % len_dataset == 0):
if opt.save_name_mode:
torch.save(
generator.module, 'MyDNN1_%s_epoch%d_bs%d.pth' %
(opt.task, epoch, opt.batch_size))
print(
'The trained model is successfully saved at epoch %d'
% (epoch))
if opt.save_mode == 'iter':
if iteration % opt.save_by_iter == 0:
if opt.save_name_mode:
torch.save(
generator.module, 'MyDNN1_%s_iter%d_bs%d.pth' %
(opt.task, iteration, opt.batch_size))
print(
'The trained model is successfully saved at iteration %d'
% (iteration))
else:
if opt.save_mode == 'epoch':
if (epoch % opt.save_by_epoch
== 0) and (iteration % len_dataset == 0):
if opt.save_name_mode:
torch.save(
generator, 'final_%s_epoch%d_bs%d.pth' %
(opt.task, epoch, opt.batch_size))
print(
'The trained model is successfully saved at epoch %d'
% (epoch))
if opt.save_mode == 'iter':
if iteration % opt.save_by_iter == 0:
if opt.save_name_mode:
torch.save(
generator, 'final_%s_iter%d_bs%d.pth' %
(opt.task, iteration, opt.batch_size))
print(
'The trained model is successfully saved at iteration %d'
% (iteration))
# ----------------------------------------
# Network dataset
# ----------------------------------------
# Define the dataloader
# trainset = dataset.TestDataset(opt)
trainset = dataset.Noise2CleanDataset(opt)
print('The overall number of training images:', len(trainset))
testset = dataset.TestDataset(opt)
valset = dataset.ValDataset(opt)
print('The overall number of val images:', len(valset))
# Define the dataloader
dataloader = DataLoader(trainset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
pin_memory=True)
val_loader = DataLoader(valset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
pin_memory=True)
test_loader = DataLoader(testset,
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers,
pin_memory=True)
# ----------------------------------------
# Training
# ----------------------------------------
# Count start time
prev_time = time.time()
best_PSNR = 0
# For loop training
for epoch in range(opt.epochs):
total_loss = 0
total_ploss = 0
total_sobel = 0
total_Lap = 0
for i, (true_input, simulated_input, true_target,
noise_level_map) in enumerate(dataloader):
# To device
true_input = true_input.cuda()
true_target = true_target.cuda()
simulated_input = simulated_input.cuda()
noise_level_map = noise_level_map.cuda()
# Train Generator
optimizer_G.zero_grad()
pre_clean = generator(true_input)
# Parse through VGGMINC layers
# features_y = vgg(pre_clean)
# features_x = vgg(true_input)
# content_loss = criterion_L2(features_y, features_x).
pre = pre_clean[0, :, :, :].data.permute(1, 2, 0).cpu().numpy()
pre = rgb2gray(pre)
true = true_input[0, :, :, :].data.permute(1, 2, 0).cpu().numpy()
true = rgb2gray(true)
laplacian_pre = cv2.Laplacian(pre, cv2.CV_32F) #CV_64F为图像深度
laplacian_gt = cv2.Laplacian(true, cv2.CV_32F) #CV_64F为图像深度
sobel_pre = 0.5 * (cv2.Sobel(pre, cv2.CV_32F, 1, 0, ksize=5) +
cv2.Sobel(pre, cv2.CV_32F, 0, 1, ksize=5)
) #1，0参数表示在x方向求一阶导数
sobel_gt = 0.5 * (cv2.Sobel(true, cv2.CV_32F, 1, 0, ksize=5) +
cv2.Sobel(true, cv2.CV_32F, 0, 1, ksize=5)
) #0,1参数表示在y方向求一阶导数
sobel_loss = mean_squared_error(sobel_pre, sobel_gt)
laplacian_loss = mean_squared_error(laplacian_pre, laplacian_gt)
# L1 Loss
Pixellevel_L1_Loss = criterion_L1(pre_clean, true_target)
# MS-SSIM loss
ms_ssim_loss = 1 - ms_ssim_module(pre_clean + 1, true_target + 1)
# Overall Loss and optimize
loss = Pixellevel_L1_Loss + 0.5 * laplacian_loss
# loss = Pixellevel_L1_Loss
loss.backward()
optimizer_G.step()
# Determine approximate time left
iters_done = epoch * len(dataloader) + i
iters_left = opt.epochs * len(dataloader) - iters_done
time_left = datetime.timedelta(seconds=iters_left *
(time.time() - prev_time))
prev_time = time.time()
total_loss = Pixellevel_L1_Loss.item() + total_loss
# total_ploss = content_loss.item() + total_ploss
total_sobel = sobel_loss + total_sobel
total_Lap = laplacian_loss + total_Lap
# # Print log
print(
"\r[Epoch %d/%d] [Batch %d/%d] [Pixellevel L1 Loss: %.4f] [laplacian_loss Loss: %.4f] [sobel_loss Loss: %.4f] Time_left: %s"
% ((epoch + 1), opt.epochs, i, len(dataloader),
Pixellevel_L1_Loss.item(), laplacian_loss.item(),
sobel_loss.item(), time_left))
img_list = [pre_clean, true_target, true_input]
name_list = ['pred', 'gt', 'noise']
utils.save_sample_png(sample_folder=sample_folder,
sample_name='MyDNN_MS_epoch%d' % (epoch + 1),
img_list=img_list,
name_list=name_list,
pixel_max_cnt=255)
# Learning rate decrease at certain epochs
lr = adjust_learning_rate(opt, (epoch + 1), (iters_done + 1),
optimizer_G)
print(
"\r[Epoch %d/%d] [Batch %d/%d] [Pixellevel L1 Loss: %.4f] [laplacian_loss Loss: %.4f] [sobel_loss Loss: %.4f] Time_left: %s"
% ((epoch + 1), opt.epochs, i, len(dataloader), total_loss / 320,
total_Lap / 320, total_sobel / 320, time_left))
### Validation
val_PSNR = 0
be_PSNR = 0
num_of_val_image = 0
for j, (true_input, simulated_input, true_target,
noise_level_map) in enumerate(val_loader):
# To device
# A is for input image, B is for target image
true_input = true_input.cuda()
true_target = true_target.cuda()
# Forward propagation
with torch.no_grad():
pre_clean = generator(true_input)
# Accumulate num of image and val_PSNR
num_of_val_image += true_input.shape[0]
val_PSNR += utils.psnr(pre_clean, true_target,
255) * true_input.shape[0]
be_PSNR += utils.psnr(true_input, true_target,
255) * true_input.shape[0]
val_PSNR = val_PSNR / num_of_val_image
be_PSNR = be_PSNR / num_of_val_image
# Record average PSNR
print('PSNR at epoch %d: %.4f' % ((epoch + 1), val_PSNR))
print('PSNR before denoising %d: %.4f' % ((epoch + 1), be_PSNR))
best_PSNR = max(val_PSNR, best_PSNR)
# Save model at certain epochs or iterations
save_model(opt, (epoch + 1), (iters_done + 1), len(dataloader),
generator, val_PSNR, best_PSNR)
