目录
- 1. tensorboard
- 1.1 本地使用
- 1.2 远程服务器使用
- 2. torchinfo
- Ref
1. tensorboard
1.1 本地使用
只需要掌握一个 torch.utils.tensorboard.writer.SummaryWriter 接口即可。
在初始化 SummaryWriter 的时候，通常需要指定log的存放路径。这个路径默认是 runs/CURRENT_DATETIME_HOSTNAME，其中
CURRENT_DATETIME_HOSTNAME = datetime.now().strftime("%b%d_%H-%M-%S") + "_" + socket.gethostname()
这里建议为每个实验单独开一个文件夹，例如 runs/exp1、runs/exp2 等，所有的events文件都会存放在 log_dir 下
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir='./runs/exp1')
writer中用到最多的接口自然是 add_scalar，接下来也将重点讲解该接口的使用。如还需参考其他接口，可以移步至官方文档。
顾名思义，add_scalar 用于逐步添加标量以形成一条曲线（例如loss、acc曲线等），具体用法如下
writer.add_scalar(tag: str, scalar_value: float, global_step: int)
tag 是该曲线的标签（例如它可以是loss、acc、lr等），scalar_value 是这一个点的纵坐标，global_step 是这一个点的横坐标。
我们可以调用 add_scalar 来绘制一条
y
=
2
x
y=2x
y=2x 的曲线，如下
import time
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('./runs/exp1')
for i in range(20):
writer.add_scalar('y=2x', 2 * i, i)
time.sleep(2)
writer.close()
上述代码每隔2秒添加一个点，40秒之后整条曲线绘制完毕。
如需查看该曲线，打开本地终端，执行命令（既可以在绘制的过程中执行该命令也可以在绘制结束后执行该命令）
tensorboard --logdir=./runs/
📝 关于
--logdir的官方解释：Directory where TensorBoard will look to find TensorFlow event files that it can display. TensorBoard will recursively walk the directory structure rooted at logdir, looking for.*tfevents.*files. A leading tilde will be expanded with the semantics of Python’sos.expanduserfunction.
然后在本地浏览器中打开 http://localhost:6006/ 即可，效果如下
我们还可以对 tag 进行分级，例如
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('./runs/exp2')
for i in range(20):
writer.add_scalar('straight_line/y=2x', 2 * i, i)
writer.add_scalar('straight_line/y=3x', 3 * i, i)
writer.add_scalar('parabola/y=x^2', i * i, i)
writer.add_scalar('parabola/y=3x^2', 3 * i * i, i)
writer.close()
效果如下
我们还可以在同一个 tag 上添加多条曲线
import random
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('./runs/exp3')
for i in range(5):
slopes = [1, 2, 3, 4, 5]
scope = random.randint(15, 40)
for j in range(scope):
writer.add_scalar('different_lines', slopes[i] * j, j)
writer.close()
效果如下
tensorboard的默认端口是6006，如果该端口事先被占用，则tensorboard无法正常启动。我们可以手动指定端口 --port=[指定的端口号] 也可以让系统自动寻找可用的端口 --port=0。
1.2 远程服务器使用
由于远程服务器上通常是没有浏览器的，因此我们需要进行相应的配置才可以在本地的浏览器访问远程服务器上的tensorboard服务。
首先建立一个SSH隧道，在本地机器上执行
ssh -NfL 6006:localhost:6006 remote_username@remote_host
-N 代表不执行远程命令，-f 代表把SSH挂在后台。该命令将本地的6006端口转发至远程服务器的6006端口。之后，在远程服务器上启动tensorboard服务
tensorboard --logdir=/path/to/logs
最后，在本地机器上打开http://localhost:6006/即可。
注意，以上方式建立的SSH隧道会一直挂在后台，如果需要关闭隧道，可执行以下命令
pkill -f 'ssh -NfL'
如果不加 -Nf，则执行命令后会直接连接到远程服务器，当断开与远程服务器的连接后，SSH隧道也会随之关闭。
2. torchinfo
通常我们会使用如下语法来计算一个模型的总参数量
print(sum(p.numel() for p in model.parameters()))
但是这种方式并不能提供关于模型各层的详细信息。对于复杂的模型，我们可能更关心每层的参数量、输出形状、以及内存占用等信息。在这种情况下，我们可以使用 torchinfo。
torchinfo 的前身是 torchsummary，安装方式如下
pip install torchinfo
同样地，只需要掌握一个 torchinfo.torchinfo.summary 接口即可。
查看一个模型的完整信息需要提供输入的形状（如果只想查看参数量则不必提供），且形状应当包含 batch_size 这个维度。例如，查看resnet18的信息需要提供形状为
(
N
,
C
,
H
,
W
)
=
(
N
,
3
,
224
,
224
)
(N,C,H,W)=(N,3,224,224)
(N,C,H,W)=(N,3,224,224) 的输入：
from torchinfo import summary
from torchvision.models import resnet18
model = resnet18()
batch_size = 32
summary(model, input_size=(batch_size, 3, 224, 224))
输出如下
==========================================================================================
Layer (type:depth-idx) Output Shape Param #
==========================================================================================
ResNet [32, 1000] --
├─Conv2d: 1-1 [32, 64, 112, 112] 9,408
├─BatchNorm2d: 1-2 [32, 64, 112, 112] 128
├─ReLU: 1-3 [32, 64, 112, 112] --
├─MaxPool2d: 1-4 [32, 64, 56, 56] --
├─Sequential: 1-5 [32, 64, 56, 56] --
│ └─BasicBlock: 2-1 [32, 64, 56, 56] --
│ │ └─Conv2d: 3-1 [32, 64, 56, 56] 36,864
│ │ └─BatchNorm2d: 3-2 [32, 64, 56, 56] 128
│ │ └─ReLU: 3-3 [32, 64, 56, 56] --
│ │ └─Conv2d: 3-4 [32, 64, 56, 56] 36,864
│ │ └─BatchNorm2d: 3-5 [32, 64, 56, 56] 128
│ │ └─ReLU: 3-6 [32, 64, 56, 56] --
│ └─BasicBlock: 2-2 [32, 64, 56, 56] --
│ │ └─Conv2d: 3-7 [32, 64, 56, 56] 36,864
│ │ └─BatchNorm2d: 3-8 [32, 64, 56, 56] 128
│ │ └─ReLU: 3-9 [32, 64, 56, 56] --
│ │ └─Conv2d: 3-10 [32, 64, 56, 56] 36,864
│ │ └─BatchNorm2d: 3-11 [32, 64, 56, 56] 128
│ │ └─ReLU: 3-12 [32, 64, 56, 56] --
├─Sequential: 1-6 [32, 128, 28, 28] --
│ └─BasicBlock: 2-3 [32, 128, 28, 28] --
│ │ └─Conv2d: 3-13 [32, 128, 28, 28] 73,728
│ │ └─BatchNorm2d: 3-14 [32, 128, 28, 28] 256
│ │ └─ReLU: 3-15 [32, 128, 28, 28] --
│ │ └─Conv2d: 3-16 [32, 128, 28, 28] 147,456
│ │ └─BatchNorm2d: 3-17 [32, 128, 28, 28] 256
│ │ └─Sequential: 3-18 [32, 128, 28, 28] 8,448
│ │ └─ReLU: 3-19 [32, 128, 28, 28] --
│ └─BasicBlock: 2-4 [32, 128, 28, 28] --
│ │ └─Conv2d: 3-20 [32, 128, 28, 28] 147,456
│ │ └─BatchNorm2d: 3-21 [32, 128, 28, 28] 256
│ │ └─ReLU: 3-22 [32, 128, 28, 28] --
│ │ └─Conv2d: 3-23 [32, 128, 28, 28] 147,456
│ │ └─BatchNorm2d: 3-24 [32, 128, 28, 28] 256
│ │ └─ReLU: 3-25 [32, 128, 28, 28] --
├─Sequential: 1-7 [32, 256, 14, 14] --
│ └─BasicBlock: 2-5 [32, 256, 14, 14] --
│ │ └─Conv2d: 3-26 [32, 256, 14, 14] 294,912
│ │ └─BatchNorm2d: 3-27 [32, 256, 14, 14] 512
│ │ └─ReLU: 3-28 [32, 256, 14, 14] --
│ │ └─Conv2d: 3-29 [32, 256, 14, 14] 589,824
│ │ └─BatchNorm2d: 3-30 [32, 256, 14, 14] 512
│ │ └─Sequential: 3-31 [32, 256, 14, 14] 33,280
│ │ └─ReLU: 3-32 [32, 256, 14, 14] --
│ └─BasicBlock: 2-6 [32, 256, 14, 14] --
│ │ └─Conv2d: 3-33 [32, 256, 14, 14] 589,824
│ │ └─BatchNorm2d: 3-34 [32, 256, 14, 14] 512
│ │ └─ReLU: 3-35 [32, 256, 14, 14] --
│ │ └─Conv2d: 3-36 [32, 256, 14, 14] 589,824
│ │ └─BatchNorm2d: 3-37 [32, 256, 14, 14] 512
│ │ └─ReLU: 3-38 [32, 256, 14, 14] --
├─Sequential: 1-8 [32, 512, 7, 7] --
│ └─BasicBlock: 2-7 [32, 512, 7, 7] --
│ │ └─Conv2d: 3-39 [32, 512, 7, 7] 1,179,648
│ │ └─BatchNorm2d: 3-40 [32, 512, 7, 7] 1,024
│ │ └─ReLU: 3-41 [32, 512, 7, 7] --
│ │ └─Conv2d: 3-42 [32, 512, 7, 7] 2,359,296
│ │ └─BatchNorm2d: 3-43 [32, 512, 7, 7] 1,024
│ │ └─Sequential: 3-44 [32, 512, 7, 7] 132,096
│ │ └─ReLU: 3-45 [32, 512, 7, 7] --
│ └─BasicBlock: 2-8 [32, 512, 7, 7] --
│ │ └─Conv2d: 3-46 [32, 512, 7, 7] 2,359,296
│ │ └─BatchNorm2d: 3-47 [32, 512, 7, 7] 1,024
│ │ └─ReLU: 3-48 [32, 512, 7, 7] --
│ │ └─Conv2d: 3-49 [32, 512, 7, 7] 2,359,296
│ │ └─BatchNorm2d: 3-50 [32, 512, 7, 7] 1,024
│ │ └─ReLU: 3-51 [32, 512, 7, 7] --
├─AdaptiveAvgPool2d: 1-9 [32, 512, 1, 1] --
├─Linear: 1-10 [32, 1000] 513,000
==========================================================================================
Total params: 11,689,512
Trainable params: 11,689,512
Non-trainable params: 0
Total mult-adds (G): 58.05
==========================================================================================
Input size (MB): 19.27
Forward/backward pass size (MB): 1271.92
Params size (MB): 46.76
Estimated Total Size (MB): 1337.94
==========================================================================================
可以看到，resnet18的大小为11M。一个浮点数占用四个字节，那么所有参数将占用 11689512 ⋅ 4 / 1 0 6 ≈ 46.76 11689512\cdot4/10^6\approx46.76 11689512⋅4/106≈46.76 MB的内存。输入占用19.27MB的内存，进行一次前向/反向传播占用1271.92MB的内存（因为要保存中间变量），所以在训练resnet18的过程中将总共占用1337.94MB的内存。
有些时候，我们不想立即让 summary 打印出模型的信息，因此可以设置 verbose=0，等到合适的时候进行打印
model_stats = summary(model, input_size=(batch_size, 3, 224, 224), verbose=0)
# ...
print(model_stats)
如果模型需要提供多个输入，我们可以为 input_size 提供 List[Tuple[int, ...]] 的格式，例如Transformer
from torchinfo import summary
import torch
model = torch.nn.Transformer()
batch_size = 128
src_len, tgt_len = 256, 256
embed_dim = 512
summary(model, input_size=[(src_len, batch_size, embed_dim), (tgt_len, batch_size, embed_dim)])
输出如下
====================================================================================================
Layer (type:depth-idx) Output Shape Param #
====================================================================================================
Transformer [256, 128, 512] --
├─TransformerEncoder: 1-1 [256, 128, 512] --
│ └─ModuleList: 2-1 -- --
│ │ └─TransformerEncoderLayer: 3-1 [256, 128, 512] 3,152,384
│ │ └─TransformerEncoderLayer: 3-2 [256, 128, 512] 3,152,384
│ │ └─TransformerEncoderLayer: 3-3 [256, 128, 512] 3,152,384
│ │ └─TransformerEncoderLayer: 3-4 [256, 128, 512] 3,152,384
│ │ └─TransformerEncoderLayer: 3-5 [256, 128, 512] 3,152,384
│ │ └─TransformerEncoderLayer: 3-6 [256, 128, 512] 3,152,384
│ └─LayerNorm: 2-2 [256, 128, 512] 1,024
├─TransformerDecoder: 1-2 [256, 128, 512] --
│ └─ModuleList: 2-3 -- --
│ │ └─TransformerDecoderLayer: 3-7 [256, 128, 512] 4,204,032
│ │ └─TransformerDecoderLayer: 3-8 [256, 128, 512] 4,204,032
│ │ └─TransformerDecoderLayer: 3-9 [256, 128, 512] 4,204,032
│ │ └─TransformerDecoderLayer: 3-10 [256, 128, 512] 4,204,032
│ │ └─TransformerDecoderLayer: 3-11 [256, 128, 512] 4,204,032
│ │ └─TransformerDecoderLayer: 3-12 [256, 128, 512] 4,204,032
│ └─LayerNorm: 2-4 [256, 128, 512] 1,024
====================================================================================================
Total params: 44,140,544
Trainable params: 44,140,544
Non-trainable params: 0
Total mult-adds (G): 6.46
====================================================================================================
Input size (MB): 134.22
Forward/backward pass size (MB): 12348.03
Params size (MB): 100.92
Estimated Total Size (MB): 12583.17
====================================================================================================
可以看到，transformer的大小为44M，训练时将占用12GB的内存。
Ref
[1] https://stackoverflow.com/questions/37987839/how-can-i-run-tensorboard-on-a-remote-server
[2] https://github.com/TylerYep/torchinfo
