def main(_):
if not os.path.exists(args.ckpt_dir):
os.makedirs(args.ckpt_dir)
if not os.path.exists(args.sample_dir):
os.makedirs(args.sample_dir)
if not os.path.exists(args.test_dir):
os.makedirs(args.test_dir)
if args.use_gpu:
# added to controll the gpu memory
print("GPU\n")
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.9)
with tf.Session(config=tf.ConfigProto(
gpu_options=gpu_options)) as sess:
model = DnCNN(sess,
sigma=args.sigma,
lr=args.lr,
dataset=args.trainset)
if args.phase == 'train':
model.train()
else:
model.test()
else:
print("CPU\n")
with tf.Session() as sess:
model = DnCNN(sess,
sigma=args.sigma,
lr=args.lr,
dataset=args.trainset)
if args.phase == 'train':
model.train()
else:
model.test()
if cuda:
decompose_model = decompose_model.cuda()
# compose_model = compose_model.cuda()
# device_ids = [0]
# model = nn.DataParallel(model, device_ids=device_ids).cuda()
# criterion = criterion.cuda()
optimizer_decompose = optim.Adam(decompose_model.parameters(), lr=args.lr)
scheduler_decompose = MultiStepLR(optimizer_decompose,
milestones=[30, 60, 90],
gamma=0.2) # learning rates
# optimizer_compose = optim.Adam(compose_model.parameters(), lr=args.lr)
# scheduler_compose = MultiStepLR(optimizer_compose, milestones=[30, 60, 90], gamma=0.2) # learning rates
for epoch in range(initial_epoch, n_epoch):
decompose_model.train()
# compose_model.train()
scheduler_decompose.step(
epoch) # step to the learning rate in this epcoh
# scheduler_compose.step(epoch) # step to the learning rate in this epcoh
xs = dg.datagenerator(data_dir=args.train_data)
xs = xs.astype('float32') / 255.0
xs = torch.from_numpy(xs.transpose(
(0, 3, 1, 2))) # tensor of the clean patches, NXCXHXW
DDataset = DenoisingDataset(xs, sigma)
batch_y, batch_x = DDataset[:238336]
dataset = torch.cat((batch_x, batch_y), dim=1)
if not os.path.exists(save_dir):
os.mkdir(save_dir)
if __name__ == '__main__':
# model selection
print('===> Building model')
model = DnCNN()
low_model = DnCNN()
low_model.load_state_dict(torch.load(os.path.join(args.low_model_dir, args.low_model_name)))
initial_epoch = findLastCheckpoint(save_dir=save_dir) # load the last model in matconvnet style
if initial_epoch > 0:
print('resuming by loading epoch %03d' % initial_epoch)
# model.load_state_dict(torch.load(os.path.join(save_dir, 'model_%03d.pth' % initial_epoch)))
# model = torch.load(os.path.join(save_dir, 'model_%03d.pth' % initial_epoch))
model.train()
low_model.eval()
# criterion = nn.MSELoss(reduction = 'sum') # PyTorch 0.4.1
# criterion = sum_squared_error()
criterion = nn.MSELoss()
Edge_enhance = torch.FloatTensor(args.batch_size, 1, 40, 40)
if cuda:
model = model.cuda()
low_model = low_model.cuda()
# device_ids = [0]
# model = nn.DataParallel(model, device_ids=device_ids).cuda()
criterion = criterion.cuda()
optimizer = optim.Adam(model.parameters(), lr=args.lr)
scheduler = MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.2) # learning rates
for epoch in range(initial_epoch, n_epoch):
def train_model(config):
# Define hyper-parameters.
depth = int(config["DnCNN"]["depth"])
n_channels = int(config["DnCNN"]["n_channels"])
img_channel = int(config["DnCNN"]["img_channel"])
kernel_size = int(config["DnCNN"]["kernel_size"])
use_bnorm = config.getboolean("DnCNN", "use_bnorm")
epochs = int(config["DnCNN"]["epoch"])
batch_size = int(config["DnCNN"]["batch_size"])
train_data_dir = config["DnCNN"]["train_data_dir"]
test_data_dir = config["DnCNN"]["test_data_dir"]
eta_min = float(config["DnCNN"]["eta_min"])
eta_max = float(config["DnCNN"]["eta_max"])
dose = float(config["DnCNN"]["dose"])
model_save_dir = config["DnCNN"]["model_save_dir"]
# Save logs to txt file.
log_dir = config["DnCNN"]["log_dir"]
log_dir = Path(log_dir) / "dose{}".format(str(int(dose * 100)))
log_file = log_dir / "train_result.txt"
# Define device.
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Initiate a DnCNN instance.
# Load the model to device and set the model to training.
model = DnCNN(depth=depth, n_channels=n_channels,
img_channel=img_channel,
use_bnorm=use_bnorm,
kernel_size=kernel_size)
model = model.to(device)
model.train()
# Define loss criterion and optimizer
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scheduler = MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.2)
criterion = LossFunc(reduction="mean")
# Get a validation test set and corrupt with noise for validation performance.
# For every epoch, use this pre-determined noisy images.
test_file_list = glob.glob(test_data_dir + "/*.png")
xs_test = []
# Can't directly convert the xs_test from list to ndarray because some images are 512*512
# while the rest are 256*256.
for i in range(len(test_file_list)):
img = cv2.imread(test_file_list[i], 0)
img = np.array(img, dtype="float32") / 255.0
img = np.expand_dims(img, axis=0)
img_noisy, _ = nm(img, eta_min, eta_max, dose, t=100)
xs_test.append((img_noisy, img))
# Train the model.
loss_store = []
epoch_loss_store = []
psnr_store = []
ssim_store = []
psnr_tr_store = []
ssim_tr_store = []
loss_mse = torch.nn.MSELoss()
dtype = torch.cuda.FloatTensor
# load vgg network
vgg = Vgg16().type(dtype)
for epoch in range(epochs):
# For each epoch, generate clean augmented patches from the training directory.
# Convert the data from uint8 to float32 then scale them to make it in [0, 1].
# Then make the patches to be of shape [N, C, H, W],
# where N is the batch size, C is the number of color channels.
# H and W are height and width of image patches.
xs = dg.datagenerator(data_dir=train_data_dir)
xs = xs.astype("float32") / 255.0
xs = torch.from_numpy(xs.transpose((0, 3, 1, 2)))
train_set = dg.DenoisingDatatset(xs, eta_min, eta_max, dose)
train_loader = DataLoader(dataset=train_set, num_workers=4,
drop_last=True, batch_size=batch_size,
shuffle=True) # TODO: if drop_last=True, the dropping in the
# TODO: data_generator is not necessary?
# train_loader_test = next(iter(train_loader))
t_start = timer()
epoch_loss = 0
for idx, data in enumerate(train_loader):
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
img_batch_read = len(inputs)
optimizer.zero_grad()
outputs = model(inputs)
# We can use labels for both style and content image
# style image
# style_transform = transforms.Compose([
# normalize_tensor_transform() # normalize with ImageNet values
# ])
# labels_t = style_transform(labels)
labels_t = labels.repeat(1, 3, 1, 1)
outputs_t = outputs.repeat(1, 3, 1, 1)
y_c_features = vgg(labels_t)
style_gram = [gram(fmap) for fmap in y_c_features]
y_hat_features = vgg(outputs_t)
y_hat_gram = [gram(fmap) for fmap in y_hat_features]
# calculate style loss
style_loss = 0.0
for j in range(4):
style_loss += loss_mse(y_hat_gram[j], style_gram[j][:img_batch_read])
style_loss = STYLE_WEIGHT*style_loss
aggregate_style_loss = style_loss
# calculate content loss (h_relu_2_2)
recon = y_c_features[1]
recon_hat = y_hat_features[1]
content_loss = CONTENT_WEIGHT*loss_mse(recon_hat, recon)
aggregate_content_loss = content_loss
loss = aggregate_content_loss + aggregate_style_loss
# loss = criterion(outputs, labels)
loss_store.append(loss.item())
epoch_loss += loss.item()
loss.backward()
optimizer.step()
if idx % 100 == 0:
print("Epoch [{} / {}], step [{} / {}], loss = {:.5f}, lr = {:.6f}, elapsed time = {:.2f}s".format(
epoch + 1, epochs, idx, len(train_loader), loss.item(), *scheduler.get_last_lr(), timer()-t_start))
epoch_loss_store.append(epoch_loss / len(train_loader))
# At each epoch validate the result.
model = model.eval()
# # Firstly validate on training sets. This takes a long time so I commented.
# tr_psnr = []
# tr_ssim = []
# # t_start = timer()
# with torch.no_grad():
# for idx, train_data in enumerate(train_loader):
# inputs, labels = train_data
# # print(inputs.shape)
# # inputs = np.expand_dims(inputs, axis=0)
# # inputs = torch.from_numpy(inputs).to(device)
# inputs = inputs.to(device)
# labels = labels.squeeze().numpy()
#
# outputs = model(inputs)
# outputs = outputs.squeeze().cpu().detach().numpy()
#
# tr_psnr.append(peak_signal_noise_ratio(labels, outputs))
# tr_ssim.append(structural_similarity(outputs, labels))
# psnr_tr_store.append(sum(tr_psnr) / len(tr_psnr))
# ssim_tr_store.append(sum(tr_ssim) / len(tr_ssim))
# # print("Elapsed time = {}".format(timer() - t_start))
#
# print("Validation on train set: epoch [{} / {}], aver PSNR = {:.2f}, aver SSIM = {:.4f}".format(
# epoch + 1, epochs, psnr_tr_store[-1], ssim_tr_store[-1]))
# Validate on test set
val_psnr = []
val_ssim = []
with torch.no_grad():
for idx, test_data in enumerate(xs_test):
inputs, labels = test_data
inputs = np.expand_dims(inputs, axis=0)
inputs = torch.from_numpy(inputs).to(device)
labels = labels.squeeze()
outputs = model(inputs)
outputs = outputs.squeeze().cpu().detach().numpy()
val_psnr.append(peak_signal_noise_ratio(labels, outputs))
val_ssim.append(structural_similarity(outputs, labels))
psnr_store.append(sum(val_psnr) / len(val_psnr))
ssim_store.append(sum(val_ssim) / len(val_ssim))
print("Validation on test set: epoch [{} / {}], aver PSNR = {:.2f}, aver SSIM = {:.4f}".format(
epoch + 1, epochs, psnr_store[-1], ssim_store[-1]))
# Set model to train mode again.
model = model.train()
scheduler.step()
# Save model
save_model(model, model_save_dir, epoch, dose * 100)
# Save the loss and validation PSNR, SSIM.
if not log_dir.exists():
Path.mkdir(log_dir)
with open(log_file, "a+") as fh:
# fh.write("{} Epoch [{} / {}], loss = {:.6f}, train PSNR = {:.2f}dB, train SSIM = {:.4f}, "
# "validation PSNR = {:.2f}dB, validation SSIM = {:.4f}".format(
# datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"),
# epoch + 1, epochs, epoch_loss_store[-1],
# psnr_tr_store[-1], ssim_tr_store[-1],
# psnr_store[-1], ssim_store[-1]))
fh.write("{} Epoch [{} / {}], loss = {:.6f}, "
"validation PSNR = {:.2f}dB, validation SSIM = {:.4f}\n".format(
datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:"),
epoch + 1, epochs, epoch_loss_store[-1],
psnr_store[-1], ssim_store[-1]))
# np.savetxt(log_file, np.hstack((epoch + 1, epoch_loss_store[-1], psnr_store[-1], ssim_store[-1])), fmt="%.6f", delimiter=", ")
fig, ax = plt.subplots()
ax.plot(loss_store[-len(train_loader):])
ax.set_title("Last 1862 losses")
ax.set_xlabel("iteration")
fig.show()