MMonitor
A simple tool for observing the internal dynamics of neural networks.
作者:DLCV Team
这是一个监测网络训练过程中指标的轻量级工具,即插即用,也可以方便的自己编写自己想要观察的指标。
工具的两个重要概念(维度)是模块和指标,模块指的是pytorch中的神经网络模块例如Linear、Conv2d,指标是指想要观察的指标例如权重范数等等。可以针对单个模块设置多个观察指标,同一个指标可以被多个模块使用(如果计算支持)。
工具的架构借鉴了软件工程中的三层架构:1.数据准备(通过注册hook获取指标计算需要的数据),2.指标计算(将1中得到的基础数据进行二次加工)3.界面(用户与工具交互的接口)。此外,工具还将指标计算与数据展示解耦,使用者可以根据自己熟悉的可视化工具将计算得到的数据进行展示。
[TOC]
1. 使用方法
1.1 安装
// 下载源代码
git clone https://openi.pcl.ac.cn/LLM-Research/ModelEvolution/src/branch/master/MMonitor
// 进入下载目录
cd 到下载目录,与setup.py文件同级
// 在安装前可以激活你想要安装的环境后运行安装代码
pip install -e ./ 或者 python setup.py develop
1.2 使用
from MMonitor.mmonitor.monitor import Monitor # 工具接口类
import wandb # 展示工具,用wandb展示(可选,计算结果可以通过调用Monitor.get_output()方法获取)
from MMonitor.visualize import Visualization # 可选(可以不用)
"""
编写想要观察模块以及对应的观察指标
基本格式:{
模块1:[[指标1,对应设置],[指标2,对应设置]],
模块2:[[指标1,对应设置],[指标2,对应设置]]
}
观察模块写法支持
1. 'fc1':model中自己定义的模块名称
2. 'Conv2d':pytorch中的模块类名称的字符串形式
3. nn.BatchNorm2d:pytorch中的模块类名称
注意: 2,3方式是模糊搜索,会将所有满足条件的模块都进行观察,例如2是模型中所有的卷积模块都观察
观察指标写法支持
注意:指标的对应设置(例如 linear(5, 2))是指每隔多少step(minibatch)计算一次指标,linear是指线性(目前仅支持linear方式),5是指每隔5个step计算一次,2是指从第2 个step开始计数也就是第2、7、12...个step的时候进行计算
1. 'MeanTID':指标的类名称的字符串形式,这种形式默认对应设置是linear(0, 0)
2. ['MeanTID']同1
3. ['MeanTID', 'linear(5, 0)']:使用MeanTID指标,这个指标计算从0开始,每隔5个计算一次
"""
config = {
nn.BatchNorm2d: [['MeanTID', 'linear(5,0)'], 'InputSndNorm']
}
model = Model()
monitor = Monitor(model, config)
step = 0
for epoch in range(epochs):
for data in dataloader:
y = model(data)
loss.backward()
###########################
monitor.track(step)
step += 1
monitor.get_output()
1.3 运行结果展示
monitor_config = {
'Linear': [['WeightNorm', 'linear(100, 0)'], ['InputSndNorm', 'linear(100, 0)'], ['OutputGradSndNorm', 'linear(100, 0)']],
}
model = vit_s()
2.支持的指标
指标分为两类一类是singlestep,一类是multistep
singlestep:是指指标计算在一个step内就可以完成的指标
multistep:是指标计算需要多个step结果聚合在一起才能计算完成的指标
2.1 Single Step Quantity
Name |
描述 |
实现 |
Extension |
cite |
InputCovMaxEig |
输入协方差矩阵的最大特征值 |
1. data = module.input.cov_matrix_eigs 2. Max Eig value |
ForwardInputEigOfCovExtension |
https://arxiv.org/pdf/2002.10801.pdf |
InputCovStableRank |
输入协方差矩阵的稳定秩 |
1. data = module.input.cov_matrix_eigs 2. Eigs sum / Max Eig |
ForwardInputEigOfCovExtension |
|
InputCovCondition20 |
输入协方差矩阵的20%条件数 |
1. data = module.input.cov_matrix_eigs 2. top20% Eig values |
ForwardInputEigOfCovExtension |
https://arxiv.org/pdf/2002.10801.pdf |
InputCovCondition50 |
输入协方差矩阵的50%条件数 |
1. data = module.input.cov_matrix_eigs 2. top50% Eig values |
ForwardInputEigOfCovExtension |
https://arxiv.org/pdf/2002.10801.pdf |
InputCovCondition80 |
输入协方差矩阵的80%条件数 |
1. data = module.input.cov_matrix_eigs 2. top80% Eig values |
ForwardInputEigOfCovExtension |
https://arxiv.org/pdf/2002.10801.pdf |
WeightNorm |
权重二范数 |
1. data = module.weight 2. norm(2) |
|
|
InputMean |
输入的每个channel的均值 |
1. data = module.input 2. mean |
ForwardInputExtension |
|
OutputGradSndNorm |
输出梯度二范数 |
1. data = module.output_grad 2. norm(2) |
BackwardOutputExtension |
|
InputSndNorm |
输出二范数 |
1. data = module.input 2. norm(2) |
ForwardInputExtension |
|
2.2 Multi Step Quantity
Name |
描述 |
实现 |
Extension |
cite |
MeanTID |
BN模块中batch的训练和推理时mean的差异 |
1.data = module.input 2.datas = [data.mean] 3.diff_data = [d-module.running_mean for d in datas] 4. stack diff_data 5. diff_data.norm(-1)/(sqrt(running_var).norm(-1)) 6. mean (思路详见论文公式) |
ForwardInputExtension |
https://arxiv.org/abs/2210.05153 |
VarTID |
BN模块中batch的训练和推理时var的差异 |
1. data = input 2. datas = [sqrt(data.var(1))] 3. sigma = sqrt(running_var) 4. diff_data = [d-sigma for d in datas] 5. stack diff_data 6. diff_data(-1)/(sigma.norm(-1)) 7. mean(思路详见论文公式) |
ForwardInputExtension |
https://arxiv.org/abs/2210.05153 |
3. 开发者
3.1 Extension
e.g. 计算InputMean指标时需要知道当前模块的输入,但是当前模块获取不到,因此需要通过注册hook先获取input,Extension的作用就是获取input的hook并将获取的input作为当前模块的一个属性。Extension根据获取方式分为forward Extension和backward Extension
3.1.1 Forward Extension
Name |
描述 |
实现 |
ForwardInputExtension |
获取模块的输入 |
1.return input[0] |
ForwardOutputExtension |
获取模块的输出 |
1.return output[0] |
ForwardInputEigOfCovExtension |
获取输入协方差矩阵的特征值 |
1.cal_cov_matrix(data) 2.retrun cal_eig(matrix) |
3.1.2 Backward Extension
Name |
描述 |
实现 |
BackwardInputExtension |
获取模块输入的梯度 |
1. return grad_input[0] |
BackwardOutputExtension |
获取模块输出的梯度 |
1. return grad_output[0] |
BackwardOutputEigOfCovExtension |
获取模块输出梯度协方差矩阵的特征值 |
1. cal_cov_matrix(grad_output) 2. return cal_eig(matrix) |
3.1.3 utils
Function |
描述 |
实现 |
cal_cov_matrix |
计算数据的协方差矩阵 |
return torch.cov(data.T) |
cal_eig |
计算数据的特征值 |
return torch.linalg.eigvals(data) |
3.2 Visualization
visualize的作用是将monitor的output提取出来,在wandb中进行展示,主要有三个函数
- show(step):将第step个数据提取出来进行展示,可以在这个函数中进行数据的存储
- log_ext():等价于wandb.log()
- close():wandb.finish()
注:monitor的output的数据组织方式
{
module_name:{
quantity:{
step:{
data
}
}
}
}