使用 TensorBoard 的 GAN 模型摘要 Pytorch?

时间:2021-01-04 07:10:19

标签: pytorch tensorboard

有没有办法可以使用 Pytorch 在 TensorBoard 中可视化 GAN 架构的完整训练循环?我认为可以使用 TF,但我很难找到使用 Pytorch 的方法。

1 个答案:

答案 0 :(得分:0)

您可以为此使用 TensorboardX

您可以使用 TensorboardX 中的 SummaryWriter 在给定目录中创建事件文件并向其中添加摘要和事件。

以下代码是您可以使用的示例,但您必须自己添加损失值、真实图像和生成的图像。我评论了他们必须去的地方。

from tensorboardX import SummaryWriter

import torchvision.utils as vutils
import numpy as np

REPORT_EVERY_ITER = 100
SAVE_IMAGE_EVERY_ITER = 1000

if __name__ == "__main__":

    writer = SummaryWriter()

    gen_losses = []
    dis_losses = []
    iter_no = 0

    // looping over the batches in the environment 
    for batch_v in iterate_batches(envs):

        // getting the outputs 
        // getting the generators loss 
        // getting the discriminators loss

        iter_no += 1

        // save the loss values for both generators and the discriminator every 100 steps 

        if iter_no % REPORT_EVERY_ITER == 0:
            log.info(
                "Iter %d: gen_loss=%.3e, dis_loss=%.3e",
                iter_no,
                np.mean(gen_losses),
                np.mean(dis_losses),
            )
            writer.add_scalar("gen_loss", np.mean(gen_losses), iter_no)
            writer.add_scalar("dis_loss", np.mean(dis_losses), iter_no)
            gen_losses = []
            dis_losses = []

        // save the images being produced from both the ground truth and the generator
        // it is saved every 1000 iterations
        if iter_no % SAVE_IMAGE_EVERY_ITER == 0:
            // save the generated images from the generator 
            writer.add_image(
                "fake",
                vutils.make_grid(gen_output_v.data[:64], normalize=True),
                iter_no
            )
            // add the ground truth images here
            // these will be the same throughout the cycle 
            writer.add_image(
                "real", 
                vutils.make_grid(batch_v.data[:64], normalize=True), 
                iter_no
            )

要查看结果,只需在运行模型训练的同一目录中运行命令:tensorboard --logdir runsruns 包含训练结果)。将显示一个链接,您可以通过该链接查看如下图所示的图。如果您想在远程服务器上运行 Tensorboard,则必须在命令行中添加命令 --bind_all 才能从外部访问它。

查看生成的图像

enter image description here

查看损失值

enter image description here