pytorch中tensorboardX进行可视化

环境依赖：

pytorch   0.4以上

tensorboardX:   pip install tensorboardX、pip install tensorflow

 

在项目代码中加入tensorboardX的记录代码，生成文件并返回到浏览器中显示可视化结果。

官方示例：

 

默认设置是在根目录下生成一个runs文件夹，里面存储summary的信息。

在runs的同级目录下命令行中输入：

tensorboard --logdir runs            (不是输tensorboardX)

会出来一个网站,复制到浏览器即可可视化loss,acc,lr等数据的变化过程.

 

举例说明pytorch中设置summary的方式：实例化summary对象，然后在对象上add_scalar想要监督的指标即可。

包含调用logging包对log文件保存方法,详见链接(https://www.cnblogs.com/ywheunji/p/14125085.html)

class Trainer(object):
    def __init__(self, weight_path, resume, gpu_id, accumulate, fp_16):
        init_seeds(0)
        self.train_dataloader = DataLoader(
            self.train_dataset,
            batch_size=cfg.TRAIN["BATCH_SIZE"],
            num_workers=cfg.TRAIN["NUMBER_WORKERS"],
            shuffle=True,
            pin_memory=True,
        )

        self.yolov4 = Build_Model(weight_path=weight_path, resume=resume).to(
            self.device
        )

    def train(self):
        global writer
        logger.info(
            "Training start,img size is: {:d},batchsize is: {:d},work number is {:d}".format(
                cfg.TRAIN["TRAIN_IMG_SIZE"],
                cfg.TRAIN["BATCH_SIZE"],
                cfg.TRAIN["NUMBER_WORKERS"],
            )
        )
        logger.info(self.yolov4)
        logger.info(
            "Train datasets number is : {}".format(len(self.train_dataset))
        )

        if self.fp_16:
            self.yolov4, self.optimizer = amp.initialize(
                self.yolov4, self.optimizer, opt_level="O1", verbosity=0
            )
        logger.info("        =======  start  training   ======     ")
        for epoch in range(self.start_epoch, self.epochs):
            start = time.time()
            self.yolov4.train()

            mloss = torch.zeros(4)
            logger.info("===Epoch:[{}/{}]===".format(epoch, self.epochs))
            for i, (imgs, label_sbbox,
            ) in enumerate(self.train_dataloader):

                loss, loss_ciou, loss_conf, loss_cls = self.criterion(p, p_d, label_sbbox)

                loss.backward()
                # Print batch results
                if i % 10 == 0:
                    logger.info(
                        "  === Epoch:[{:3}/{}],step:[{:3}/{}],img_size:[{:3}],total_loss:{:.4f}|loss_ciou:{:.4f}|loss_conf:{:.4f}|loss_cls:{:.4f}|lr:{:.4f}".format(
                            epoch,
                            self.epochs,
                            i,
                            len(self.train_dataloader) - 1,
                            self.train_dataset.img_size,
                            mloss[3],
                            mloss[0],
                            mloss[1],
                            mloss[2],
                            self.optimizer.param_groups[0]["lr"],
                        )
                    )
                    writer.add_scalar(
                        "loss_ciou",
                        mloss[0],
                        len(self.train_dataloader)
                        * epoch
                        + i,
                    )
                    writer.add_scalar(
                        "train_loss",
                        mloss[3],
                        len(self.train_dataloader)
                        * epoch
                        + i,
                    )


            # eval
            logger.info(
                "===== Validate =====".format(epoch, self.epochs)
            )
            logger.info("val img size is {}".format(cfg.VAL["TEST_IMG_SIZE"]))
            with torch.no_grad():
                APs, inference_time = Evaluator(
                    self.yolov4, showatt=False
                ).APs_voc()
                for i in APs:
                    logger.info("{} --> mAP : {}".format(i, APs[i]))
                    mAP += APs[i]
                mAP = mAP / self.train_dataset.num_classes
                logger.info("mAP : {}".format(mAP))
                logger.info(
                    "inference time: {:.2f} ms".format(inference_time)
                )
                writer.add_scalar("mAP", mAP, epoch)
                self.__save_model_weights(epoch, mAP)
                logger.info("save weights done")
            logger.info("  ===test mAP:{:.3f}".format(mAP))
            
if __name__ == "__main__":
    global logger, writer
    writer = SummaryWriter(logdir=opt.log_path + "/event")
    logger = Logger(
        log_file_name=opt.log_path + "/log.txt",
        log_level=logging.DEBUG,
        logger_name="YOLOv4",
    ).get_log()

    Trainer(
        weight_path=opt.weight_path,
        resume=opt.resume,
        gpu_id=opt.gpu_id,
        accumulate=opt.accumulate,
        fp_16=opt.fp_16,
    ).train()