1.在训练代码中创建一个SummaryWriter对象,指定存储路径和命名空间。
例如:
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir='logs/mobilenetv2', comment='MobileNetV2')
其中,log_dir参数指定TensorBoard日志的存储路径,comment参数指定命名空间,可以用来区分不同的实验结果。
2.在训练过程中,将需要记录的信息写入SummaryWriter对象。例如,记录每个epoch的损失值和准确率:
for epoch in range(num_epochs):
# train the model
train_loss, train_acc = train(...)
# validate the model
val_loss, val_acc = validate(...)
# write to TensorBoard
writer.add_scalar('Loss/train', train_loss, epoch)
writer.add_scalar('Loss/val', val_loss, epoch)
writer.add_scalar('Accuracy/train', train_acc, epoch)
writer.add_scalar('Accuracy/val', val_acc, epoch)
其中,add_scalar方法用来写入一个标量数据,第一个参数指定数据的名称,第二个参数指定数据的值,第三个参数指定数据所在的step或epoch。
3.记录模型结构和参数信息。可以使用SummaryWriter对象的add_graph方法来记录模型的计算图,使用add_histogram方法来记录参数的分布情况。
例如:
# record model structure
writer.add_graph(model, input_to_model=torch.randn(1, 3, 224, 224))
# record parameter distributions
for name, param in model.named_parameters():
writer.add_histogram(name, param, epoch)
其中,add_graph方法的第一个参数是模型对象,第二个参数是一个输入张量,用来生成计算图。add_histogram方法的第一个参数是数据的名称,第二个参数是参数张量,第三个参数是数据所在的step或epoch。
4.在训练结束后,关闭SummaryWriter对象并保存TensorBoard日志。
例如:
writer.close()
运行以上代码,就可以在TensorBoard中查看网络的模型结构、训练过程和参数变化情况了。需要注意的是,在TensorBoard中查看参数分布情况时,可以选择使用Histograms或Distributions来展示,两者的区别是前者是将参数值分成若干个区间并统计每个区间内参数的数量,后者是将参数值视为一个连续的概率分布,并计算其概率密度函数。