caffe模型转 pytorch 模型

矫情吗；* 2022-11-10 14:15 279阅读 0赞

**最近基于 caff2onnx 做了部分修改，完成了caffe 转 pytorch的模型代码，**

主代码 , 需要自己构建 pytorch 的Net 架构， 同时 net各层的名字要与 caffe的各层对应。

graph, params = LoadCaffeModel(caffe_graph_path,caffe_params_path)
        #print(graph)
    
        net_pytorch = NET()
    
        net_pytorch.eval()
    
    
        print('start convert')
        caffe2pytorch = Caffe2Pytorch(graph, params, net_pytorch)
        caffe2pytorch.convert()
        print("convert finish.")
        torch.save(net_pytorch.state_dict(), save_path)
        print("save finish.")

子函数代码

class Caffe2Pytorch():
        def __init__(self, net, model, pytorch_net):
            # 初始化一个c2oGraph对象
            # 网络和参数
            self.netLayerCaffe = self.GetNetLayerCaffe(net)
            self.netModelCaffe = self.GetNetModelCaffe(model)
    
            # 模型的输入名和输入维度
            self.model_input_name = []
            self.model_input_shape = []
    
            self.pytorch_net = pytorch_net
            self.state_dict = {}
    
        # 获取网络层
        def GetNetLayerCaffe(self, net):
            if len(net.layer) == 0 and len(net.layers) != 0:
                return net.layers
            elif len(net.layer) != 0 and len(net.layers) == 0:
                return net.layer
            else:
                print("prototxt layer error")
                return -1
    
        # 获取参数层
        def GetNetModelCaffe(self, model):
            if len(model.layer) == 0 and len(model.layers) != 0:
                return model.layers
            elif len(model.layer) != 0 and len(model.layers) == 0:
                return model.layer
            else:
                print("caffemodel layer error")
                return -1
    
    
        def match(self, caffe_layer_name, pS, pD):
            index = 0
            for name in self.pytorch_net.state_dict():
                if name.find(caffe_layer_name+".") == 0:
                    print("match success:  caffe name:", caffe_layer_name,  " py name:", name)
                    newD = [p for p in pD[index]]
                    newD = np.array(newD)
                    shape = tuple([s for s in pS[index]])
                    newD = newD.reshape(shape)
                    #print("newD:", newD.shape)
                    self.state_dict[name] = torch.from_numpy(newD)
                    index += 1
    
                if index == len(pS):
                    break
    
    
        def convert(self):
            ParamShape = []
            ParamData = []
            # 根据这个layer名找出对应的caffemodel中的参数
            for i, model_layer in enumerate(self.netModelCaffe):
                Params = copy.deepcopy(model_layer.blobs)
                ParamShape = [p.shape.dim for p in Params]
                ParamData = [p.data for p in Params]
    
                if model_layer.type == "BatchNorm" or model_layer.type == "BN":
                    if len(ParamShape) == 3:
                        # 如果是bn层，params为[mean, var, s]，则需要把mean和var除以滑动系数s
                        ParamShape = ParamShape[:-1]
                        ParamData = [
                            [q / (Params[-1].data[0])
                             for q in p.data] if i == 0 else
                            [q / (Params[-1].data[0] + 1e-5) for q in p.data]
                            for i, p in enumerate(Params[:-1])
                        ]  # with s
                    elif len(ParamShape) == 2 and len(ParamShape[0]) == 4:
                        ParamShape = [[ParamShape[0][1]], [ParamShape[1][1]]]
                        ParamData = [[q / 1. for q in p.data] if i == 0 else
                                     [q / (1. + 1e-5) for q in p.data]
                                     for i, p in enumerate(Params)]
                    if self.netModelCaffe[i+1].type == "Scale":
                        Params = copy.deepcopy(self.netModelCaffe[i+1].blobs)
                        ParamShape1 = [p.shape.dim for p in Params]
                        ParamData1 = [p.data for p in Params]
    
                        ParamShape1.extend(ParamShape)
                        ParamData1.extend(ParamData)
                        ParamShape = ParamShape1
                        ParamData = ParamData1
    
                print("caffe param name:", model_layer.name, " param shape :", ParamShape)
    
                layer_name = model_layer.name
                layer_name = layer_name.replace("/", "_")
                self.match(layer_name, ParamShape, ParamData)
    
            self.pytorch_net.load_state_dict(self.state_dict, strict=False)
    
            return ParamShape, ParamData