def test(model_path, model):
device = torch.device('cpu')
model = Flow(in_channel, 64, n_block, n_flows, img_sz, num_scales)
checkpoint = torch.load(model_path, map_location=device)
model.load_state_dict(checkpoint)
z_sample = torch.randn((16, 3, 32, 32), dtype=torch.float32)
x = model.reverse(z_sample)
x = torch.sigmoid(x)
show(x[0].permute(1, 2, 0).detach().cpu(), "test", 0)
optimizer.zero_grad()
x = x.to(device)
# show(x[0].detach().permute(1, 2, 0), os.path.join(root, "input/1.png"), 0)
z, sldj = model(x)
# x = model.reverse(z)
# show(x[0].detach().permute(1, 2, 0), os.path.join(root, "input/1.png"), 0)
loss = loss_fn(z, sldj)
loss.backward()
grad_clipping(optimizer)
optimizer.step()
if loss.item() < old_loss:
is_best = 1
save_checkpoint(model.state_dict(), is_best, root, f"./models/checkpoint.pth.tar")
old_loss = loss
print(f"loss at batch {b} epoch {epoch}: {loss.item()}")
if b == 0: #len(train_loader) - 1:
z_sample = torch.randn((16, 3, 32, 32), dtype=torch.float32).to(device)
model.eval()
if use_cuda:
x = model.module.reverse(z_sample)
else:
x = model.reverse(z_sample)
images = torch.sigmoid(x)
images_concat = utils.make_grid(images, nrow=int(16 ** 0.5), padding=2, pad_value=255)
utils.save_image(images_concat, os.path.join(root, f"./samples/{epoch}_{b}_{str((torch.round(loss * 10**3)/10**3).item())}.png"))
# show(x.squeeze().detach().cpu(), os.path.join(root, f"./samples/{epoch}_{b}_{str((torch.round(loss * 10**3)/10**3).item())}.png"), 1)
# show(x.squeeze().permute(1, 2, 0).detach().cpu(), os.path.join(root, f"./samples/{epoch}_{b}_{str((torch.round(loss * 10**3)/10**3).item())}.png"), 1)
best_loss = old_loss
print(f"Best loss at {best_loss}")