在训练和评估模型的过程中,我们需要保存很多的数据,比如loss,precision、recall、f1等等。通过这些信息,来更好地了解模型的情况,对模型进行优化。
一般来说,我们用pytorch的tensorboard就够了。但是我发现tensorboard有几个缺点:
- 需要本地开一个服务器:
tensorboard --logdir=/xxx/log --port 6008 --bind_all
- 每次重新运行的时候都得关掉服务器再重新运行才能看到新的图
在看了师兄的代码以后,发现了这么一个工具: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
})