本帖最后由 丙丁先生 于 2024-7-5 07:32 编辑
浮点-定点两阶段量化训练方案介绍1. RawMode图中的RawMode可以是已经训练完成的浮点模型,也可以是随机初始化的模型。 RawModel是已经训练完成的浮点模型:
- 可以认为FloatStage 进行浮点数据调整,其中最重要的调整包括ConvBn融合,参数normalize和激活值normalize,这种调整有利于Fix Stage 的定点训练
- 在实际使用中,由已训练的 raw model 直接进行FixStage 对定点损失比较严重,甚至loss直接崩溃
RawModel是随机初始化的模型:
- 可以认为FloatStage 的策略会直接参与训练,包括对weight和激活的normalize策略,以及ConvBn融合
- 此处要注意ConvBN融合,并非简单直接将BN参数推入Conv,而是前向融合,反向不融合策略,获得最大的训练收益
- 可能对训练效率有少许影响,经过几个案例的测试,我们发现影响大概在10%内
2. RawStateDict图中RawStateDict加载是向后兼容的,兼容顺序可以箭头方向 - 在模型变换的每个节点(图中FloatModel和FloatModelWithFixScale),都可以进行对StateDict加载和保存
- 同一训练节点Save的StateDict可以被本节点Model加载(Save到Load的左向箭头)
- 可以被后续阶段的节点加载
- 在每个阶段可以导出onnx
- 注意: 模型加载需要在网络变换的最后加载,中间不允许加载模型,如下代码示意model=MyNet()model=linger.trace_layers(....)model.load_state_dict(...) #it's OKCopy
model=MyNet()model=linger.trace_layers(model,....)#model.load_state_dict(...) #it's NOT OKmodel=linger.normalize_layers(model,....)model.load_state_dict(...) #it's OKCopy
model=MyNet()model=linger.trace_layers(model,....)#model.load_state_dict(...) #it's NOT OKmodel=linger.normalize_layers(model,....)#model.load_state_dict(...) #it's NOT OKmodel=linger.init(model,...)model.load_state_dict(...) #it's OKCopy
3. layers 相关策略3.1 ConvBn融合原理介绍linger 浮点阶段实现了convbn融合,执行保守融合策略 使用实例如下 class Model(nn.Module): def __init__(self): super(Model, self).__init__() self.transpose = nn.ConvTranspose2d(10, 10, 5, 5, 2, 4, 2, True, 2) self.conv = nn.Conv2d(10, 10, kernel_size=3, stride=1, padding=1, bias=True, groups=2) self.bn = nn.BatchNorm2d(10) self.fc = nn.Linear(10*254*254, 100) def forward(self, x): x = self.transpose(x) x = self.conv(x) x = self.bn(x) n, c, h, w = x.shape x = x.view(n, c*h*w) x = self.fc(x) return xmodel = Model()print(model)Copy
Model( (transpose): ConvTranspose2d(10, 10, kernel_size=(5, 5), stride=(5, 5), padding=(2, 2), dilation=(2, 2), output_padding=(4, 4), groups=2) (conv): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2) (bn): BatchNorm2d(10, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (fc): Linear(in_features=645160, out_features=100, bias=True))Copy
dummy_input = torch.randn(1, 10, 10, 10).cuda()linger.trace_layers(model, model, dummy_input, fuse_bn=True)print(model)Copy
Model( (transpose): ConvTranspose2d(10, 10, kernel_size=(5, 5), stride=(5, 5), padding=(2, 2), dilation=(2, 2), output_padding=(4, 4), groups=2) (conv): NormalizeConvBN2d( normalize_data:None,normalize_weight:None,normalize_bias:None,ahead_relu:False (conv): Conv2d(10, 10, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=2) (bn): BatchNorm2d(10, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) ) (bn): EmptyBatchNorm() (fc): Linear(in_features=645160, out_features=100, bias=True))Copy
从两次的输出的net网络可以看出,lingerCONVBN融合策略会将原网络的bn替换成EmptyBatchNorm,将原来的conv替换成NormalizeConvBN2d。 - EmptyBatchNorm : 什么都没做,仅仅占位
- NormalizeConvBN2d : 带有normalize能力的convbn2d模块(此处normalize为None,不设置normalize功能)
3.2 AHEAD_RELU众所周知,Relu 操作是将负值用零替换。但如果Relu前面的操作OPX输出的量化,对正负值有一定的偏向性,特别是在负向有较大幅值,非常不利于输OPX的量化。由于OPX后紧跟Relu,因此可以使用OPX量化时可以仅仅关注正值量化,此策略即为AHEAD_RELU。
例如如下网络定义 class Model(nn.Module): def __init__(self): super(Model, self).__init__() self.conv = nn.Conv2d(10, 10, kernel_size=3, stride=1, padding=1, bias=True) self.bn = nn.BatchNorm2d(10) self.relu = nn.ReLU() self.conv1 = nn.Conv2d(10, 10, kernel_size=3, stride=1, padding=1, bias=True) self.relu1 = nn.ReLU() self.conv2 = nn.Conv2d(10, 10, kernel_size=3, stride=1, padding=1, bias=True) self.bn1 = nn.BatchNorm2d(10) self.relu2 = nn.ReLU() self.fc = nn.Linear(250, 100) def forward(self, x): x = self.conv(x) x = self.bn(x) x = self.relu(x) x = x - 1 x = self.conv1(x) x = self.relu1(x) x = self.conv2(x) x = self.bn1(x) x = self.relu2(x) n, c, h, w = x.shape x = x.view((n, c*h*w)) x = self.fc(x) return xCopy
其中 self.bn, self.conv1, self.b1都可以仅仅关注正方向的量化。 当前支持的组合策略包括 ahead_conv_relu,ahead_bn_relu,ahead_linear_relu
在trace_layers接口中均有选项,用户可以选用
trace_layers(root_model,target_model,args,*,fuse_bn:bool =True,ahead_conv_relu:bool =True,ahead_bn_relu:bool =True,ahead_linear_relu:bool =True) 注意: - 虽然trace_layers ahead relu是属于浮点(预)训练阶段,但仅仅为配置选项,真正使用和产生影响是在量化训练阶段,即init之后
- trace_layers接口需要在init之前调用,同时只能执行一次,目前不允许网络trace多个部分, 建议放在所有配置操作的最前面
|