调用 ofa_model()
以获取 OFA 模型。此方法会查找所有 nn.Conv2d
/ nn.ConvTranspose2d
和 nn.BatchNorm2d
模块,并将这些模块替换为 DynamicConv2d
/ DynamicConvTranspose2d
和 DynamicBatchNorm2d
。
在最终 OFA 模型中,需包含剪枝率列表以便为每个层指定最大和最小通道剪枝率。
OFA 模型中每个卷积层的输出通道均可使用任意的剪枝率。此列表中的最大值和最小值分别表示该模型的最大和最小压缩率。列表中的其他值则表示要最优化的子网络。剪枝率默认设为 [0.5, 0.75, 1]。
对于从 OFA 模型中采样的子网络,卷积层的输出通道是剪枝率列表中的数字之一乘以其原始数字,例如,对于剪枝率列表 [0.5, 0.75, 1] 和卷积层 nn.Conv2d(16, 32, 5)
,该层在采样子网络中的输出通道是 [0.5*32, 0.75*32, 1*32] 之一。
由于第一层和最后一层对网络性能存在显著影响,因此通常排除在剪枝之外。默认情况下,此方法会自动识别第一个卷积和最后一个卷积,并将其置于排除列表中。将 auto_add_excludes
设置为 False 即可取消此功能特性。
ofa_model = ofa_pruner.ofa_model([0.5, 0.75, 1], excludes = None, auto_add_excludes=True)