criterion = cnn.loss
optimizer = optim.Adam(cnn.parameters(), weight_decay=.001)
# optimizer = optim.SGD(cnn.parameters(), lr=1e-2, momentum=0.9)
logger = Logger('./log')
logger.add_graph(cnn, image_size=96)
index = 0
for epoch in range(400):
for i, data in enumerate(trainloader):
inputs, labels = data
lightness, z_truth, original = inputs['lightness'], inputs[
'z_truth'], inputs['original_lab_image']
optimizer.zero_grad()
outputs = cnn(lightness)
ab_outputs = cnn.decode_ab_values()
colorized_im = torch.cat((lightness, ab_outputs), 1)
# plot_image_channels(colorized_im.detach()[0, :, :, :], figure=20)
loss = criterion(z_truth)
loss.backward()
optimizer.step()
# Logging loss to tensorboardx
info = {'loss': loss}
for tag, value in info.items():
print(value.detach())
logger.scalar_summary(tag, value.detach(), epoch)
# Displaying Zhat
logger.histogram_summary('Zhat', outputs.detach(), epoch)
