深度学习可视化工具wandb

在训练和评估模型的过程中，我们需要保存很多的数据，比如loss，precision、recall、f1等等。通过这些信息，来更好地了解模型的情况，对模型进行优化。

一般来说，我们用pytorch的tensorboard就够了。但是我发现tensorboard有几个缺点：

1. 需要本地开一个服务器：

   tensorboard --logdir=/xxx/log --port 6008 --bind_all
   

2. 每次重新运行的时候都得关掉服务器再重新运行才能看到新的图

在看了师兄的代码以后，发现了这么一个工具：wandb，全称是Weights & Biases。

wandb的优点在于，你记录的信息会保存到云端：https://wandb.ai/。而且，每次运行，都会在云端自动创建一个新的"run"。并且，信息也比tensorboard更丰富、界面更好看！颜控无法抗拒啊！

下面是我使用时的一些截图：

---

下面讲一下怎么用吧。

首先，在本地的shell安装wandb:

pip install wandb

在wandb官网注册，当然你可以直接使用github登录。之后，在setting中，找到你的API key：

然后在本地的shell中：

wandb login 你的API key

在项目中使用也只需要简单的几步即可。具体可以参考wandb quickstart

import wandb  # 导入

wandb.init(project="my-test-project")  # 初始化

wandb.config.batch_size = 32  # 记录超参数

wandb.log({'accuracy': train_acc, 'loss': train_loss})  # 记录训练过程中的数据

---

下面是我写的Logger类。其中，config是一个保存配置的类。

# import
import wandb
from torch.utils.tensorboard import SummaryWriter


class Config(object):
    def __init__(self, logger_name="wandb"):
        self.project_name = "casee"
        self.logger_name = logger_name
        self.log_dir = "./data/log/"
		... ...
        
        
class Logger(object):
    def __init__(self, config):
        self.config = config
        self.logger_name = self.config.logger_name
        self.logger = self._get_logger()

    def _get_logger(self):
        logger = None

        if self.logger_name == "wandb":
            wandb.init(project=self.config.project_name,
                       config=self.config.__dict__)
            logger = wandb

        elif self.logger_name == "tensorboard":
            logger = SummaryWriter(self.config.log_dir)

        return logger

    def log(self, info):
        if self.logger_name == "wandb":
            self.logger.log(info)
        elif self.logger_name == "tensorboard":
            main_tag, tag_scalar_dict, global_step = info
            self.logger.add_scalars(main_tag, tag_scalar_dict, global_step)

调用示例：

logger = Logger(config)
logger.log({
	"Train/loss": train_loss
})