def build_model(self):
# RCAN model
self.output = self.SRDenseNet(
x=self.x_lr,
f=self.n_filters,
kernel_size=self.kernel_size,
reduction=self.reduction,
use_bn=self.use_bn,
scale=self.img_scale,
)
self.output = tf.clip_by_value(self.output * 255., 0., 255.)
# l1 loss
# self.loss = tf.reduce_mean(tf.abs(self.output-self.x_hr))
# l2 loss
self.loss = tf.losses.mean_squared_error(self.x_hr, self.output)
self.train_op = self.opt.minimize(self.loss,
global_step=self.global_step)
# metrics
self.psnr = tf.reduce_mean(
metric.psnr(self.x_hr, self.output, m_val=255))
self.ssim = tf.reduce_mean(
metric.ssim(self.x_hr, self.output, m_val=255))
# summaries
tf.summary.image('lr', self.x_lr, max_outputs=self.batch_size)
tf.summary.image('hr', self.x_hr, max_outputs=self.batch_size)
tf.summary.image('generated-hr',
self.output,
max_outputs=self.batch_size)
tf.summary.scalar("loss/l2_loss", self.loss)
tf.summary.scalar("metric/psnr", self.psnr)
tf.summary.scalar("metric/ssim", self.ssim)
tf.summary.scalar("misc/lr", self.lr)
# merge summary
self.merged = tf.summary.merge_all()
# model saver
self.saver = tf.train.Saver(max_to_keep=2)
self.best_saver = tf.train.Saver(max_to_keep=1)
def build_model(self):
# RCAN model
self.output = self.residual_channel_attention_network(
x=self.x_lr,
f=self.n_filters,
kernel_size=self.kernel_size,
reduction=self.reduction,
use_bn=self.use_bn,
scale=self.img_scale,
)
self.output = tf.clip_by_value(self.output * 255., 0., 255.)
# l1 loss
self.loss = tf.reduce_mean(tf.abs(self.output - self.x_hr))
self.train_op = self.opt.minimize(self.loss,
global_step=self.global_step)
# metrics
self.psnr = tf.reduce_mean(metric.psnr(self.output, self.x_hr,
m_val=1))
self.ssim = tf.reduce_mean(metric.ssim(self.output, self.x_hr,
m_val=1))
# summaries
tf.summary.image('lr', self.x_lr, max_outputs=self.batch_size)
tf.summary.image('hr', self.x_hr, max_outputs=self.batch_size)
tf.summary.image('generated-hr',
self.output,
max_outputs=self.batch_size)
tf.summary.scalar("loss/l1_loss", self.loss)
tf.summary.scalar("metric/psnr", self.psnr)
tf.summary.scalar("metric/ssim", self.ssim)
tf.summary.scalar("misc/lr", self.lr)
# merge summary
self.merged = tf.summary.merge_all()
# model saver
self.saver = tf.train.Saver(max_to_keep=1)
self.best_saver = tf.train.Saver(max_to_keep=1)
self.writer = tf.summary.FileWriter(self.tf_log, self.sess.graph)
if torch.cuda.is_available():
torch.set_default_tensor_type("torch.cuda.FloatTensor")
torch.backends.cudnn.benchmark = True
MAX_EPOCHS = 10
resize = Resize((128, 128))
dataset = Dataset("datasets/test/test12", resize)
loader = DataLoader(dataset, batch_size=4)
model = PReNet_r(recurrent_iter=6).to(device)
optimizer = Adam(model.parameters(), lr=0.001)
scheduler = MultiStepLR(optimizer, milestones=[30, 50, 80], gamma=0.2)
for epoch in range(MAX_EPOCHS):
scheduler.step()
for x, y in loader:
x = x.to(device)
y = y.to(device)
y_pred = model(x)
loss = -ssim(y, y_pred)
loss.backward()
with torch.no_grad():
mse = torch.mean((y - y_pred)**2)
p = psnr(mse)
s = -loss
print(loss.item(), p.item())
optimizer.step()
optimizer.zero_grad()
body = tfu.conv2d(x, f=f, k=kernel_size, name="conv2d-body")
body += head # tf.math.add(body, head)
# 3. tail
x = self.up_scaling(body, f, scale, name='up-scaling')
tail = tfutil.conv2d(x, f=self.n_channel, k=kernel_size, name="conv2d-tail") # (-1, 384, 384, 3)
tail = tfu.conv2d(x, f=self.n_channel, k=kernel_size, name="conv2d-tail") # (-1, 384, 384, 3)
x = self.image_processing(tail, sign=1, name='post-processing')
return x
@@ -236,8 +237,8 @@ def build_model(self):
self.train_op = self.opt.minimize(self.loss, global_step=self.global_step)
# metrics
self.psnr = tf.reduce_mean(metric.psnr(self.output, self.x_hr, m_val=1.))
self.ssim = tf.reduce_mean(metric.ssim(self.output, self.x_hr, m_val=1.))
self.psnr = tf.reduce_mean(metric.psnr(self.output, self.x_hr, m_val=1))
self.ssim = tf.reduce_mean(metric.ssim(self.output, self.x_hr, m_val=1))
# summaries
tf.summary.image('lr', self.x_lr, max_outputs=self.batch_size)
tf.summary.image('hr', self.x_hr, max_outputs=self.batch_size)
tf.summary.image('generated-hr', self.output, max_outputs=self.batch_size)
tf.summary.scalar("loss/l1_loss", self.loss)
tf.summary.scalar("metric/psnr", self.psnr)
tf.summary.scalar("metric/ssim", self.ssim)
tf.summary.scalar("misc/lr", self.lr)
# merge summary
self.merged = tf.summary.merge_all()
# model saver
self.saver = tf.train.Saver(max_to_keep=1)
net.load_state_dict(torch.load('weights/unet_final.pt', map_location=torch.device('cpu')))
# net = IM2HI()
# net.load_state_dict(torch.load('weights/im2height_augment.pt', map_location=torch.device('cpu')))
net.eval()
data = testloader()
data = iter(data)
for i in range(3):
img, mask = data.next()
with torch.set_grad_enabled(False):
output = net(img)
print(ssim(output, mask))
print('l1', torch.nn.L1Loss()(output, mask))
print('l2', torch.nn.MSELoss()(output, mask))
# print('none l1', torch.nn.L1Loss(reduction='none')(output, mask))
# print('none l2', torch.nn.MSELoss(reduction='none')(output, mask))
output = torch.squeeze(output, 0)
output = torch.cat((output, output, output))
img, mask = torch.squeeze(img, 0), torch.squeeze(mask, 0)
mask = torch.cat((mask, mask, mask))
fig, (a1, a2,a3) = plt.subplots(1, 3, figsize=(15,5))
a1.imshow(img.permute(1,2,0))
a2.imshow(mask.permute(1,2,0))
a3.imshow(output.permute(1,2,0))