摘要：本文旨在分享 Pytorch->Caffe->om 模型转换流程。
标准网络
Baseline：PytorchToCaffe
主要功能代码在：
PytorchToCaffe+-- Caffe| +-- caffe.proto| +-- layer_param.py+-- example| +-- resnet_pytorch_2_caffe.py+-- pytorch_to_caffe.py
复制代码
直接使用可以参考resnet_pytorch_2_caffe.py，如果网络中的操作 Baseline 中都已经实现，则可以直接转换到 Caffe 模型。
添加自定义操作
如果遇到没有实现的操作，则要分为两种情况来考虑。
Caffe 中有对应操作
以 arg_max 为例分享一下添加操作的方式。
首先要查看 Caffe 中对应层的参数：caffe.proto为对应版本 caffe 层与参数的定义，可以看到 ArgMax 定义了out_max_val、top_k、axis三个参数：
message ArgMaxParameter { // If true produce pairs (argmax, maxval) optional bool out_max_val = 1 [default = false]; optional uint32 top_k = 2 [default = 1]; // The axis along which to maximise -- may be negative to index from the // end (e.g., -1 for the last axis). // By default ArgMaxLayer maximizes over the flattened trailing dimensions // for each index of the first / num dimension. optional int32 axis = 3;}
复制代码
与Caffe算子边界中的参数是一致的。
layer_param.py构建了具体转换时参数类的实例，实现了操作参数从 Pytorch 到 Caffe 的传递：
def argmax_param(self, out_max_val=None, top_k=None, dim=1): argmax_param = pb.ArgMaxParameter() if out_max_val is not None: argmax_param.out_max_val = out_max_val if top_k is not None: argmax_param.top_k = top_k if dim is not None: argmax_param.axis = dim self.param.argmax_param.CopyFrom(argmax_param)
复制代码
pytorch_to_caffe.py中定义了 Rp 类，用来实现 Pytorch 操作到 Caffe 操作的变换：
class Rp(object): def __init__(self, raw, replace, **kwargs): self.obj = replace self.raw = raw def __call__(self, *args, **kwargs): if not NET_INITTED: return self.raw(*args, **kwargs) for stack in traceback.walk_stack(None): if 'self' in stack[0].f_locals: layer = stack[0].f_locals['self'] if layer in layer_names: log.pytorch_layer_name = layer_names[layer] print('984', layer_names[layer]) break out = self.obj(self.raw, *args, **kwargs) return out
复制代码
在添加操作时，要使用 Rp 类替换操作：
torch.argmax = Rp(torch.argmax, torch_argmax)
复制代码
接下来，要具体实现该操作：
def torch_argmax(raw, input, dim=1): x = raw(input, dim=dim) layer_name = log.add_layer(name='argmax') top_blobs = log.add_blobs([x], name='argmax_blob'.format(type)) layer = caffe_net.Layer_param(name=layer_name, type='ArgMax', bottom=[log.blobs(input)], top=top_blobs) layer.argmax_param(dim=dim) log.cnet.add_layer(layer) return x
复制代码
即实现了 argmax 操作 Pytorch 到 Caffe 的转换。
Caffe 中无直接对应操作
如果要转换的操作在 Caffe 中无直接对应的层实现，解决思路主要有两个：
在 Pytorch 中将不支持的操作分解为支持的操作：
class InstanceNormalization(nn.Module): def __init__(self, dim, eps=1e-5): super(InstanceNormalization, self).__init__() self.gamma = nn.Parameter(torch.FloatTensor(dim)) self.beta = nn.Parameter(torch.FloatTensor(dim)) self.eps = eps self._reset_parameters() def _reset_parameters(self): self.gamma.data.uniform_() self.beta.data.zero_() def __call__(self, x): n = x.size(2) * x.size(3) t = x.view(x.size(0), x.size(1), n) mean = torch.mean(t, 2).unsqueeze(2).unsqueeze(3).expand_as(x) var = torch.var(t, 2).unsqueeze(2).unsqueeze(3).expand_as(x) gamma_broadcast = self.gamma.unsqueeze(1).unsqueeze(1).unsqueeze(0).expand_as(x) beta_broadcast = self.beta.unsqueeze(1).unsqueeze(1).unsqueeze(0).expand_as(x) out = (x - mean) / torch.sqrt(var + self.eps) out = out * gamma_broadcast + beta_broadcast return out
复制代码
但在验证HiLens Caffe算子边界中发现，om 模型转换不支持Channle维度之外的求和或求均值操作，为了规避这个操作，我们可以通过支持的算子重新实现nn.InstanceNorm2d：
class InstanceNormalization(nn.Module): def __init__(self, dim, eps=1e-5): super(InstanceNormalization, self).__init__() self.gamma = torch.FloatTensor(dim) self.beta = torch.FloatTensor(dim) self.eps = eps self.adavg = nn.AdaptiveAvgPool2d(1) def forward(self, x): n, c, h, w = x.shape mean = nn.Upsample(scale_factor=h)(self.adavg(x)) var = nn.Upsample(scale_factor=h)(self.adavg((x - mean).pow(2))) gamma_broadcast = self.gamma.unsqueeze(1).unsqueeze(1).unsqueeze(0).expand_as(x) beta_broadcast = self.beta.unsqueeze(1).unsqueeze(1).unsqueeze(0).expand_as(x) out = (x - mean) / torch.sqrt(var + self.eps) out = out * gamma_broadcast + beta_broadcast return out
复制代码
经过验证，与原操作等价，可以转为 Caffe 模型
2.在 Caffe 中通过利用现有操作实现：
在 Pytorch 转 Caffe 的过程中发现，如果存在featuremap + 6这种涉及到常数的操作，转换过程中会出现找不到 blob 的问题。我们首先查看pytorch_to_caffe.py中 add 操作的具体转换方法：
def _add(input, *args): x = raw__add__(input, *args) if not NET_INITTED: return x layer_name = log.add_layer(name='add') top_blobs = log.add_blobs([x], name='add_blob') if log.blobs(args[0]) == None: log.add_blobs([args[0]], name='extra_blob') else: layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', bottom=[log.blobs(input),log.blobs(args[0])], top=top_blobs) layer.param.eltwise_param.operation = 1 # sum is 1 log.cnet.add_layer(layer) return x
复制代码
可以看到对于 blob 不存在的情况进行了判断，我们只需要在log.blobs(args[0]) == None条件下进行修改，一个自然的想法是利用Scale层实现 add 操作：
def _add(input, *args): x = raw__add__(input, *args) if not NET_INITTED: return x layer_name = log.add_layer(name='add') top_blobs = log.add_blobs([x], name='add_blob') if log.blobs(args[0]) == None: layer = caffe_net.Layer_param(name=layer_name, type='Scale', bottom=[log.blobs(input)], top=top_blobs) layer.param.scale_param.bias_term = True weight = torch.ones((input.shape[1])) bias = torch.tensor(args[0]).squeeze().expand_as(weight) layer.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy()) log.cnet.add_layer(layer) else: layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', bottom=[log.blobs(input), log.blobs(args[0])], top=top_blobs) layer.param.eltwise_param.operation = 1 # sum is 1 log.cnet.add_layer(layer) return x
复制代码
类似的，featuremap * 6 这种简单乘法也可以通过同样的方法实现。
踩过的坑
Pooling：Pytorch 默认 ceil_mode=false，Caffe 默认 ceil_mode=true，可能会导致维度变化，如果出现尺寸不匹配的问题可以检查一下 Pooling 参数是否正确。另外，虽然文档上没有看到，但是 kernel_size > 32 后模型虽然可以转换，但推理会报错，这时可以分两层进行 Pooling 操作。
Upsample ：om 边界算子中的 Upsample 层scale_factor参数必须是 int，不能是 size。如果已有模型参数为 size 也会正常跑完 Pytorch 转 Caffe 的流程，但此时 Upsample 参数是空的。参数为 size 的情况可以考虑转为 scale_factor 或用 Deconvolution 来实现。
Transpose2d：Pytorch 中 output_padding 参数会加在输出的大小上，但 Caffe 不会，输出特征图相对会变小，此时反卷积之后的 featuremap 会变大一点，可以通过 Crop 层进行裁剪，使其大小与 Pytorch 对应层一致。另外，om 中反卷积推理速度较慢，最好是不要使用，可以用 Upsample+Convolution 替代。
Pad：Pytorch 中 Pad 操作很多样，但 Caffe 中只能进行 H 与 W 维度上的对称 pad，如果 Pytorch 网络中有h = F.pad(x, (1, 2, 1, 2), "constant", 0)这种不对称的 pad 操作，解决思路为：
如果不对称 pad 的层不存在后续的维度不匹配的问题，可以先判断一下 pad 对结果的影响，一些任务受 pad 的影响很小，那么就不需要修改。
如果存在维度不匹配的问题，可以考虑按照较大的参数充分 pad 之后进行 Crop，或是将前后两个(0, 0, 1, 1)与(1, 1, 0, 0)的 pad 合为一个(1, 1, 1, 1)，这要看具体的网络结构确定。
如果是 Channel 维度上的 pad 如F.pad(x, (0, 0, 0, 0, 0, channel_pad), "constant", 0)，可以考虑零卷积后 cat 到 featuremap 上：
zero = nn.Conv2d(in_channels, self.channel_pad, kernel_size=3, padding=1, bias=False)nn.init.constant(self.zero.weight, 0)pad_tensor = zero(x)x = torch.cat([x, pad_tensor], dim=1)
复制代码
本文分享自华为云社区《Pytorch->Caffe 模型转换》，原文作者：杜甫盖房子 。
点击关注，第一时间了解华为云新鲜技术~
评论