total_loss = 0
for i, image in enumerate(train_dataloader):
optimizer.zero_grad()
image = image.to(device)
image_recon, loss = model(image)
log_str = "---- [Epoch %d/%d, Step %d/%d] loss: %.6f----" % (
epoch, epochs, i, len(train_dataloader), loss.item())
logging.info(log_str)
loss.backward()
optimizer.step()
scheduler.step(epoch)
torch.save(model.state_dict(), f"checkpoints/ae_ckpt_%d_%.6f.pth" % (epoch, loss.item()))
model.eval()
with torch.no_grad():
z = torch.randn((1, 512)).to(device)
gen_img, _ = model.decoder(z)
gen_img = gen_img.permute(0, 2, 3, 1)
gen_img = (gen_img[0].cpu().numpy() + 1) / 2 * 255
gen_img = gen_img.astype(np.uint8)
plt.imshow(gen_img)
# plt.savefig(f"output/ae_ckpt_%d_%.6f.png" % (epoch, total_loss))
plt.show()
for i, image in enumerate(train_dataloader):
optimizer.zero_grad()
image = image.to(device)
image_recon, loss = model(image)
log_str = "---- [Epoch %d/%d, Step %d/%d] loss: %.6f----" % (
epoch, epochs, i, len(train_dataloader), loss.item())
logging.info(log_str)
loss.backward()
optimizer.step()
scheduler.step(epoch)
torch.save(model.state_dict(),
f"checkpoints/ae_ckpt_%d_%.6f.pth" % (epoch, loss.item()))
model.eval()
with torch.no_grad():
z = torch.randn((1, 512)).to(device)
gen_img, _ = model.decoder(z)
gen_img = gen_img.permute(0, 2, 3, 1)
gen_img = gen_img[0].cpu().numpy() * 255
gen_img = gen_img.astype(np.uint8)
plt.imshow(gen_img)
# plt.savefig(f"output/ae_ckpt_%d_%.6f.png" % (epoch, total_loss))
plt.show()