一、pth2caffe

Net:定义的网络结构，pth存储的模型路径

  1. def pth2caffe(Net,pth):
  2.     from pytorch2caffe import pytorch2caffe
  3.     name = 'yolov3_2'
  4.     model = Net
  5.     #将权重加载到网络中
  6.     model_dict = model.state_dict()
  7.     pretrained_dict = torch.load(pth, map_location=torch.device('cpu'))
  8.     pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) ==  np.shape(v)}
  9.     model_dict.update(pretrained_dict)
  10.     model.load_state_dict(model_dict)
  11.     model.eval() 
  12.     #定义转换格式后的输入数据大小
  13.     dummy_input = torch.ones([1, 3, 40, 40])
  15.     pytorch2caffe.trans_net(model, dummy_input, name)
  16.     pytorch2caffe.save_prototxt('{}.prototxt'.format(name))
  17.     pytorch2caffe.save_caffemodel('{}.caffemodel'.format(name))

二、测试模型生成

  1. import numpy as np
  2. import torch
  3. import torch.nn as nn
  5. #定义网络结构
  6. class YoloBody(nn.Module):
  7.     def __init__(self, anchors_mask, num_classes):
  8.         super(YoloBody, self).__init__()
  10.         self.last_layer0            = nn.Conv2d(3, len(anchors_mask[0]) * (num_classes + 5), kernel_size=1, stride=1)
  12.         self.last_layer1            = nn.Conv2d(3, len(anchors_mask[0]) * (num_classes + 5), kernel_size=1, stride=1)
  14.         self.last_layer2            = nn.Conv2d(3, len(anchors_mask[0]) * (num_classes + 5), kernel_size=1, stride=1)
  16.     def forward(self, x):
  17.         out0        = self.last_layer0(x)
  18.         out1        = self.last_layer1(x)
  19.         out2        = self.last_layer2(x)
  20.         return out0, out1, out2
  21.  #网络参数初始化       
  22. def weights_init(net, init_type='kaiming', init_gain = 0.02):
  23.     def init_func(m):
  24.         classname = m.__class__.__name__
  25.         if hasattr(m, 'weight') and classname.find('Conv') != -1:
  26.             if init_type == 'normal':
  27.                 torch.nn.init.normal_(m.weight.data, 0.0, init_gain)
  28.             elif init_type == 'xavier':
  29.                 torch.nn.init.xavier_normal_(m.weight.data, gain=init_gain)
  30.             elif init_type == 'kaiming':
  31.                 torch.nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
  32.             elif init_type == 'orthogonal':
  33.                 torch.nn.init.orthogonal_(m.weight.data, gain=init_gain)
  34.             else:
  35.                 raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
  36.         elif classname.find('BatchNorm2d') != -1:
  37.             torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
  38.             torch.nn.init.constant_(m.bias.data, 0.0)
  39.     print('initialize network with %s type' % init_type)
  40.     net.apply(init_func)
  42. def pth2caffe(Net,pth):
  43.     from pytorch2caffe import pytorch2caffe
  44.     name = 'yolov3_2'
  45.     model = Net
  46.     model_dict = model.state_dict()
  47.     pretrained_dict = torch.load(pth, map_location=torch.device('cpu'))
  48.     pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) ==  np.shape(v)}
  49.     model_dict.update(pretrained_dict)
  50.     model.load_state_dict(model_dict)
  51.     model.eval() 
  52.     #转换格式后的输入
  53.     dummy_input = torch.ones([1, 3, 40, 40])
  55.     pytorch2caffe.trans_net(model, dummy_input, name)
  56.     pytorch2caffe.save_prototxt('{}.prototxt'.format(name))
  57.     pytorch2caffe.save_caffemodel('{}.caffemodel'.format(name)) 
  58. if __name__ == "__main__":
  59.     anchors_mask    = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
  60.     #网络结构定义
  61.     Net = YoloBody(anchors_mask,20)
  62.     #自定义模型路径
  63.     modelpath = "xx/xx/yolo.pth"
  64.     weights_init(Net)
  65.     Net.train()
  66.     #保存pytorch模型
  67.     torch.save(Net.state_dict(),modelpath)
  68.     #转换格式
  69.     pth2caffe(Net,modelpath)