├── .gitignore ├── 001763.jpg ├── Caffe ├── ReadMe.md ├── __init__.py ├── caffe.proto ├── caffe_lmdb.py ├── caffe_net.py ├── caffe_pb2.py ├── layer_param.py └── net.py ├── LICENSE.md ├── Pytorch ├── __init__.py ├── augmentations.py ├── eval.py ├── train.py └── utils.py ├── README.md ├── __init__.py ├── analysis ├── CaffeA.py ├── PytorchA.py ├── README.md ├── __init__.py ├── __pycache__ │ ├── CaffeA.cpython-36.pyc │ ├── PytorchA.cpython-36.pyc │ ├── __init__.cpython-36.pyc │ ├── blob.cpython-36.pyc │ ├── layers.cpython-36.pyc │ ├── roi.cpython-36.pyc │ └── utils.cpython-36.pyc ├── blob.py ├── layers.py ├── roi.py └── utils.py ├── caffe_analyser.py ├── example ├── MobileNetV2.py ├── Resnet-annalyze.py ├── __pycache__ │ ├── MGN_analysis_example.cpython-36.pyc │ └── option.cpython-36.pyc ├── alexnet_pytorch_to_caffe.py ├── inceptionv3_pytorch_to_caffe.py ├── mobilenetV2_pytorch_to_caffe.py ├── resnet_pytorch_2_caffe.py ├── resnet_pytorch_analysis_example.py ├── verify_deploy.py └── vgg19_pytorch_to_caffe.py ├── funcs.py ├── model ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-35.pyc │ ├── __init__.cpython-36.pyc │ ├── mgn.cpython-35.pyc │ └── mgn.cpython-36.pyc ├── mgn.py └── resnet.py ├── pytorch_analyser.py └── pytorch_to_caffe.py /.gitignore: -------------------------------------------------------------------------------- 1 | 2 | \.idea/ 3 | 4 | \.vscode/ 5 | 6 | __pycache__/ 7 | 8 | *.caffemodel 9 | 10 | *.prototxt 11 | 12 | *.pth 13 | 14 | -------------------------------------------------------------------------------- /001763.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/001763.jpg -------------------------------------------------------------------------------- /Caffe/ReadMe.md: -------------------------------------------------------------------------------- 1 | # The Caffe in PytorchToCaffe Provides some convenient API 2 | If there are some problem in parse your prototxt or caffemodel, Please replace 3 | the caffe.proto with your own version and compile it with command 4 | `protoc --python_out ./ caffe.proto` 5 | 6 | ## caffe_net.py 7 | Using `from nn_tools.Caffe import caffe_net` to import this model 8 | ### Prototxt 9 | + `net=caffe_net.Prototxt(file_name)` to open a prototxt file 10 | + `net.init_caffemodel(caffe_cmd_path='caffe')` to generate a caffemodel file in the current work directory \ 11 | if your `caffe` cmd not in the $PATH, specify your caffe cmd path by the `caffe_cmd_path` kwargs. 12 | ### Caffemodel 13 | + `net=caffe_net.Caffemodel(file_name)` to open a caffemodel 14 | + `net.save_prototxt(path)` to save the caffemodel to a prototxt file (not containing the weight data) 15 | + `net.get_layer_data(layer_name)` return the numpy ndarray data of the layer 16 | + `net.set_layer_date(layer_name, datas)` specify the data of one layer in the caffemodel .`datas` is normally a list of numpy ndarray `[weights,bias]` 17 | + `net.save(path)` save the changed caffemodel 18 | ### Functions for both Prototxt and Caffemodel 19 | + `net.add_layer(layer_params,before='',after='')` add a new layer with `Layer_Param` object 20 | + `net.remove_layer_by_name(layer_name)` 21 | + `net.get_layer_by_name(layer_name)` or `net.layer(layer_name)` get the raw Layer object defined in caffe_pb2 22 | -------------------------------------------------------------------------------- /Caffe/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/Caffe/__init__.py -------------------------------------------------------------------------------- /Caffe/caffe_lmdb.py: -------------------------------------------------------------------------------- 1 | import lmdb 2 | from Caffe import caffe_pb2 as pb2 3 | import numpy as np 4 | 5 | class Read_Caffe_LMDB(): 6 | def __init__(self,path,dtype=np.uint8): 7 | 8 | self.env=lmdb.open(path, readonly=True) 9 | self.dtype=dtype 10 | self.txn=self.env.begin() 11 | self.cursor=self.txn.cursor() 12 | 13 | @staticmethod 14 | def to_numpy(value,dtype=np.uint8): 15 | datum = pb2.Datum() 16 | datum.ParseFromString(value) 17 | flat_x = np.fromstring(datum.data, dtype=dtype) 18 | data = flat_x.reshape(datum.channels, datum.height, datum.width) 19 | label=flat_x = datum.label 20 | return data,label 21 | 22 | def iterator(self): 23 | while True: 24 | key,value=self.cursor.key(),self.cursor.value() 25 | yield self.to_numpy(value,self.dtype) 26 | if not self.cursor.next(): 27 | return 28 | 29 | def __iter__(self): 30 | self.cursor.first() 31 | it = self.iterator() 32 | return it 33 | 34 | def __len__(self): 35 | return int(self.env.stat()['entries']) 36 | -------------------------------------------------------------------------------- /Caffe/caffe_net.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import caffe_pb2 as pb 3 | import google.protobuf.text_format as text_format 4 | import numpy as np 5 | from .layer_param import Layer_param 6 | 7 | class _Net(object): 8 | def __init__(self): 9 | self.net=pb.NetParameter() 10 | self.needChange = {} 11 | 12 | def layer_index(self,layer_name): 13 | # find a layer's index by name. if the layer was found, return the layer position in the net, else return -1. 14 | for i, layer in enumerate(self.net.layer): 15 | if layer.name == layer_name: 16 | return i 17 | 18 | def add_layer(self,layer_params,before='',after=''): 19 | # find the before of after layer's position 20 | index = -1 21 | if after != '': 22 | index = self.layer_index(after) + 1 23 | if before != '': 24 | index = self.layer_index(before) 25 | new_layer = pb.LayerParameter() 26 | new_layer.CopyFrom(layer_params.param) 27 | #insert the layer into the layer protolist 28 | if index != -1: 29 | self.net.layer.add() 30 | for i in range(len(self.net.layer) - 1, index, -1): 31 | self.net.layer[i].CopyFrom(self.net.layer[i - 1]) 32 | self.net.layer[index].CopyFrom(new_layer) 33 | else: 34 | self.net.layer.extend([new_layer]) 35 | 36 | def remove_layer_by_name(self,layer_name): 37 | for i,layer in enumerate(self.net.layer): 38 | if layer.name == layer_name: 39 | del self.net.layer[i] 40 | return 41 | raise(AttributeError, "cannot found layer %s" % str(layer_name)) 42 | 43 | 44 | 45 | def remove_layer_by_type(self,type_name): 46 | for i,layer in enumerate(self.net.layer): 47 | if layer.type == type_name: 48 | # self.change_layer_bottom(layer.top,layer.bottom) 49 | s1 = "\"" + layer.top[0] + "\"" 50 | s2 = "\"" + layer.bottom[0] + "\"" 51 | self.needChange[s1]=s2 52 | del self.net.layer[i] 53 | return 54 | 55 | 56 | 57 | def get_layer_by_name(self, layer_name): 58 | # get the layer by layer_name 59 | for layer in self.net.layer: 60 | if layer.name == layer_name: 61 | return layer 62 | raise(AttributeError, "cannot found layer %s" % str(layer_name)) 63 | 64 | def save_prototxt(self,path): 65 | prototxt=pb.NetParameter() 66 | prototxt.CopyFrom(self.net) 67 | for layer in prototxt.layer: 68 | del layer.blobs[:] 69 | with open(path,'w') as f: 70 | string = text_format.MessageToString(prototxt) 71 | for origin_name in self.needChange.keys(): 72 | string = string.replace(origin_name,self.needChange[origin_name]) 73 | f.write(string) 74 | 75 | def layer(self,layer_name): 76 | return self.get_layer_by_name(layer_name) 77 | 78 | def layers(self): 79 | return list(self.net.layer) 80 | 81 | 82 | 83 | class Prototxt(_Net): 84 | def __init__(self,file_name=''): 85 | super(Prototxt,self).__init__() 86 | self.file_name=file_name 87 | if file_name!='': 88 | f = open(file_name,'r') 89 | text_format.Parse(f.read(), self.net) 90 | pass 91 | 92 | def init_caffemodel(self,caffe_cmd_path='caffe'): 93 | """ 94 | :param caffe_cmd_path: The shell command of caffe, normally at /build/tools/caffe 95 | """ 96 | s=pb.SolverParameter() 97 | s.train_net=self.file_name 98 | s.max_iter=0 99 | s.base_lr=1 100 | s.solver_mode = pb.SolverParameter.CPU 101 | s.snapshot_prefix='./nn' 102 | with open('/tmp/nn_tools_solver.prototxt','w') as f: 103 | f.write(str(s)) 104 | import os 105 | os.system('%s train --solver /tmp/nn_tools_solver.prototxt'%caffe_cmd_path) 106 | 107 | class Caffemodel(_Net): 108 | def __init__(self, file_name=''): 109 | super(Caffemodel,self).__init__() 110 | # caffe_model dir 111 | if file_name!='': 112 | f = open(file_name,'rb') 113 | self.net.ParseFromString(f.read()) 114 | f.close() 115 | 116 | def save(self, path): 117 | with open(path,'wb') as f: 118 | f.write(self.net.SerializeToString()) 119 | 120 | def add_layer_with_data(self,layer_params,datas, before='', after=''): 121 | """ 122 | Args: 123 | layer_params:A Layer_Param object 124 | datas:a fixed dimension numpy object list 125 | after: put the layer after a specified layer 126 | before: put the layer before a specified layer 127 | """ 128 | self.add_layer(layer_params,before,after) 129 | new_layer =self.layer(layer_params.name) 130 | 131 | #process blobs 132 | del new_layer.blobs[:] 133 | for data in datas: 134 | new_blob=new_layer.blobs.add() 135 | for dim in data.shape: 136 | new_blob.shape.dim.append(dim) 137 | new_blob.data.extend(data.flatten().astype(float)) 138 | 139 | def get_layer_data(self,layer_name): 140 | layer=self.layer(layer_name) 141 | datas=[] 142 | for blob in layer.blobs: 143 | shape=list(blob.shape.dim) 144 | data=np.array(blob.data).reshape(shape) 145 | datas.append(data) 146 | return datas 147 | 148 | def set_layer_data(self,layer_name,datas): 149 | # datas is normally a list of [weights,bias] 150 | layer=self.layer(layer_name) 151 | for blob,data in zip(layer.blobs,datas): 152 | blob.data[:]=data.flatten() 153 | pass 154 | 155 | class Net(): 156 | def __init__(self,*args,**kwargs): 157 | raise(TypeError,'the class Net is no longer used, please use Caffemodel or Prototxt instead') 158 | -------------------------------------------------------------------------------- /Caffe/layer_param.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from . import caffe_pb2 as pb 3 | import numpy as np 4 | 5 | def pair_process(item,strict_one=True): 6 | if hasattr(item,'__iter__'): 7 | for i in item: 8 | if i!=item[0]: 9 | if strict_one: 10 | raise ValueError("number in item {} must be the same".format(item)) 11 | else: 12 | print("IMPORTANT WARNING: number in item {} must be the same".format(item)) 13 | return item[0] 14 | return item 15 | 16 | def pair_reduce(item): 17 | if hasattr(item,'__iter__'): 18 | for i in item: 19 | if i!=item[0]: 20 | return item 21 | return [item[0]] 22 | return [item] 23 | 24 | class Layer_param(): 25 | def __init__(self,name='',type='',top=(),bottom=()): 26 | self.param=pb.LayerParameter() 27 | self.name=self.param.name=name 28 | self.type=self.param.type=type 29 | 30 | self.top=self.param.top 31 | self.top.extend(top) 32 | self.bottom=self.param.bottom 33 | self.bottom.extend(bottom) 34 | 35 | def fc_param(self, num_output, weight_filler='xavier', bias_filler='constant',has_bias=True): 36 | if self.type != 'InnerProduct': 37 | raise TypeError('the layer type must be InnerProduct if you want set fc param') 38 | fc_param = pb.InnerProductParameter() 39 | fc_param.num_output = num_output 40 | fc_param.weight_filler.type = weight_filler 41 | fc_param.bias_term = has_bias 42 | if has_bias: 43 | fc_param.bias_filler.type = bias_filler 44 | self.param.inner_product_param.CopyFrom(fc_param) 45 | 46 | def conv_param(self, num_output, kernel_size, stride=(1), pad=(0,), 47 | weight_filler_type='xavier', bias_filler_type='constant', 48 | bias_term=True, dilation=None,groups=None): 49 | """ 50 | add a conv_param layer if you spec the layer type "Convolution" 51 | Args: 52 | num_output: a int 53 | kernel_size: int list 54 | stride: a int list 55 | weight_filler_type: the weight filer type 56 | bias_filler_type: the bias filler type 57 | Returns: 58 | """ 59 | if self.type not in ['Convolution','Deconvolution']: 60 | raise TypeError('the layer type must be Convolution or Deconvolution if you want set conv param') 61 | conv_param=pb.ConvolutionParameter() 62 | conv_param.num_output=num_output 63 | conv_param.kernel_size.extend(pair_reduce(kernel_size)) 64 | conv_param.stride.extend(pair_reduce(stride)) 65 | conv_param.pad.extend(pair_reduce(pad)) 66 | conv_param.bias_term=bias_term 67 | conv_param.weight_filler.type=weight_filler_type 68 | if bias_term: 69 | conv_param.bias_filler.type = bias_filler_type 70 | if dilation: 71 | conv_param.dilation.extend(pair_reduce(dilation)) 72 | if groups: 73 | conv_param.group=groups 74 | if groups != 1: 75 | conv_param.engine = 1 76 | self.param.convolution_param.CopyFrom(conv_param) 77 | 78 | def norm_param(self, eps): 79 | """ 80 | add a conv_param layer if you spec the layer type "Convolution" 81 | Args: 82 | num_output: a int 83 | kernel_size: int list 84 | stride: a int list 85 | weight_filler_type: the weight filer type 86 | bias_filler_type: the bias filler type 87 | Returns: 88 | """ 89 | l2norm_param = pb.NormalizeParameter() 90 | l2norm_param.across_spatial = False 91 | l2norm_param.channel_shared = False 92 | l2norm_param.eps = eps 93 | self.param.norm_param.CopyFrom(l2norm_param) 94 | 95 | 96 | def permute_param(self, order1, order2, order3, order4): 97 | """ 98 | add a conv_param layer if you spec the layer type "Convolution" 99 | Args: 100 | num_output: a int 101 | kernel_size: int list 102 | stride: a int list 103 | weight_filler_type: the weight filer type 104 | bias_filler_type: the bias filler type 105 | Returns: 106 | """ 107 | permute_param = pb.PermuteParameter() 108 | permute_param.order.extend([order1, order2, order3, order4]) 109 | 110 | self.param.permute_param.CopyFrom(permute_param) 111 | 112 | 113 | def pool_param(self,type='MAX',kernel_size=2,stride=2,pad=None, ceil_mode = True): 114 | pool_param=pb.PoolingParameter() 115 | pool_param.pool=pool_param.PoolMethod.Value(type) 116 | pool_param.kernel_size=pair_process(kernel_size) 117 | pool_param.stride=pair_process(stride) 118 | pool_param.ceil_mode=ceil_mode 119 | if pad: 120 | if isinstance(pad,tuple): 121 | pool_param.pad_h = pad[0] 122 | pool_param.pad_w = pad[1] 123 | else: 124 | pool_param.pad=pad 125 | self.param.pooling_param.CopyFrom(pool_param) 126 | 127 | def batch_norm_param(self,use_global_stats=0,moving_average_fraction=None,eps=None): 128 | bn_param=pb.BatchNormParameter() 129 | bn_param.use_global_stats=use_global_stats 130 | if moving_average_fraction: 131 | bn_param.moving_average_fraction=moving_average_fraction 132 | if eps: 133 | bn_param.eps = eps 134 | self.param.batch_norm_param.CopyFrom(bn_param) 135 | 136 | # layer 137 | # { 138 | # name: "upsample_layer" 139 | # type: "Upsample" 140 | # bottom: "some_input_feature_map" 141 | # bottom: "some_input_pool_index" 142 | # top: "some_output" 143 | # upsample_param { 144 | # upsample_h: 224 145 | # upsample_w: 224 146 | # } 147 | # } 148 | def upsample_param(self,size=None, scale_factor=None): 149 | upsample_param=pb.UpsampleParameter() 150 | if scale_factor: 151 | if isinstance(scale_factor,int): 152 | upsample_param.scale = scale_factor 153 | else: 154 | upsample_param.scale_h = scale_factor[0] 155 | upsample_param.scale_w = scale_factor[1] 156 | 157 | if size: 158 | if isinstance(size,int): 159 | upsample_param.upsample_h = size 160 | else: 161 | upsample_param.upsample_h = size[0] * scale_factor 162 | upsample_param.\ 163 | upsample_w = size[1] * scale_factor 164 | self.param.upsample_param.CopyFrom(upsample_param) 165 | 166 | def add_data(self,*args): 167 | """Args are data numpy array 168 | """ 169 | del self.param.blobs[:] 170 | for data in args: 171 | new_blob = self.param.blobs.add() 172 | for dim in data.shape: 173 | new_blob.shape.dim.append(dim) 174 | new_blob.data.extend(data.flatten().astype(float)) 175 | 176 | def set_params_by_dict(self,dic): 177 | pass 178 | 179 | def copy_from(self,layer_param): 180 | pass 181 | 182 | def set_enum(param,key,value): 183 | setattr(param,key,param.Value(value)) -------------------------------------------------------------------------------- /Caffe/net.py: -------------------------------------------------------------------------------- 1 | raise ImportError,'the nn_tools.Caffe.net is no longer used, please use nn_tools.Caffe.caffe_net' -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | 2 | The MIT License (MIT) 3 | 4 | Copyright (c) 2018 5 | 6 | Permission is hereby granted, free of charge, to any person obtaining a copy 7 | of this software and associated documentation files (the "Software"), to deal 8 | in the Software without restriction, including without limitation the rights 9 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 10 | copies of the Software, and to permit persons to whom the Software is 11 | furnished to do so, subject to the following conditions: 12 | 13 | The above copyright notice and this permission notice shall be included in all 14 | copies or substantial portions of the Software. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 17 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 18 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 19 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 20 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 21 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 22 | SOFTWARE. 23 | -------------------------------------------------------------------------------- /Pytorch/__init__.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | -------------------------------------------------------------------------------- /Pytorch/augmentations.py: -------------------------------------------------------------------------------- 1 | # driven from https://github.com/amdegroot/ssd.pytorch 2 | import torch 3 | from torchvision import transforms 4 | import cv2 5 | import numpy as np 6 | import types 7 | from numpy import random 8 | 9 | def intersect(box_a, box_b): 10 | max_xy = np.minimum(box_a[:, 2:], box_b[2:]) 11 | min_xy = np.maximum(box_a[:, :2], box_b[:2]) 12 | inter = np.clip((max_xy - min_xy), a_min=0, a_max=np.inf) 13 | return inter[:, 0] * inter[:, 1] 14 | 15 | 16 | def jaccard_numpy(box_a, box_b): 17 | """Compute the jaccard overlap of two sets of boxes. The jaccard overlap 18 | is simply the intersection over union of two boxes. 19 | E.g.: 20 | A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B) 21 | The box should be [x1,y1,x2,y2] 22 | Args: 23 | box_a: Single numpy bounding box, Shape: [4] or Multiple bounding boxes, Shape: [num_boxes,4] 24 | box_b: Single numpy bounding box, Shape: [4] 25 | Return: 26 | jaccard overlap: Shape: [box_a.shape[0], box_a.shape[1]] 27 | """ 28 | if box_a.ndim==1: 29 | box_a=box_a.reshape([1,-1]) 30 | inter = intersect(box_a, box_b) 31 | area_a = ((box_a[:, 2]-box_a[:, 0]) * 32 | (box_a[:, 3]-box_a[:, 1])) # [A,B] 33 | area_b = ((box_b[2]-box_b[0]) * 34 | (box_b[3]-box_b[1])) # [A,B] 35 | union = area_a + area_b - inter 36 | return inter / union # [A,B] 37 | 38 | 39 | class SSD_Compose(object): 40 | """Composes several augmentations together. 41 | Args: 42 | transforms (List[Transform]): list of transforms to compose. 43 | Example: 44 | # >>> augmentations.Compose([ 45 | # >>> transforms.CenterCrop(10), 46 | # >>> transforms.ToTensor(), 47 | # >>> ]) 48 | """ 49 | 50 | def __init__(self, transforms): 51 | self.transforms = transforms 52 | 53 | def __call__(self, img, boxes=None, labels=None): 54 | for t in self.transforms: 55 | img, boxes, labels = t(img, boxes, labels) 56 | return img, boxes, labels 57 | 58 | 59 | class SSD_Lambda(object): 60 | """Applies a lambda as a transform.""" 61 | 62 | def __init__(self, lambd): 63 | assert isinstance(lambd, types.LambdaType) 64 | self.lambd = lambd 65 | 66 | def __call__(self, img, boxes=None, labels=None): 67 | return self.lambd(img, boxes, labels) 68 | 69 | class ConvertFromInts(object): 70 | def __call__(self, image, *args): 71 | image=image.astype(np.float32) 72 | if len(args): 73 | return (image, *args) 74 | else: 75 | return image 76 | 77 | 78 | class SubtractMeans(object): 79 | def __init__(self, mean): 80 | self.mean = np.array(mean, dtype=np.float32) 81 | 82 | def __call__(self, image, *args): 83 | image = image.astype(np.float32) 84 | image -= self.mean 85 | image=image.astype(np.float32) 86 | if len(args): 87 | return (image,) 88 | else: 89 | return image 90 | 91 | 92 | class SSD_ToAbsoluteCoords(object): 93 | def __call__(self, image, boxes=None, labels=None): 94 | height, width, channels = image.shape 95 | boxes[:, 0] *= width 96 | boxes[:, 2] *= width 97 | boxes[:, 1] *= height 98 | boxes[:, 3] *= height 99 | 100 | return image, boxes, labels 101 | 102 | 103 | class SSD_ToPercentCoords(object): 104 | def __call__(self, image, boxes=None, labels=None): 105 | height, width, channels = image.shape 106 | boxes[:, 0] /= width 107 | boxes[:, 2] /= width 108 | boxes[:, 1] /= height 109 | boxes[:, 3] /= height 110 | 111 | return image, boxes, labels 112 | 113 | 114 | class BGR_2_HSV(object): 115 | def __call__(self, image, *args): 116 | HSV_img = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) 117 | if len(args): 118 | return (HSV_img, *args) 119 | else: 120 | return HSV_img 121 | 122 | class HSV_2_BGR(object): 123 | def __call__(self, image, *args): 124 | BGR_img = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) 125 | if len(args): 126 | return (BGR_img, *args) 127 | else: 128 | return BGR_img 129 | 130 | class Resize(object): 131 | def __init__(self, size=300): 132 | self.size = size 133 | 134 | def __call__(self, image, *args): 135 | image = cv2.resize(image, (self.size, 136 | self.size)) 137 | if len(args): 138 | return (image,) 139 | else: 140 | return image 141 | 142 | 143 | class RandomSaturation(object): 144 | def __init__(self, lower=0.5, upper=1.5): 145 | self.lower = lower 146 | self.upper = upper 147 | assert self.upper >= self.lower, "contrast upper must be >= lower." 148 | assert self.lower >= 0, "contrast lower must be non-negative." 149 | 150 | def __call__(self, image, *args): 151 | if random.randint(2): 152 | image[:, :, 1] = (image[:,:,1]*random.uniform(self.lower, self.upper)).astype(np.uint8) 153 | if len(args): 154 | return (image, *args) 155 | else: 156 | return image 157 | 158 | class RandomHue(object): 159 | def __init__(self, delta=15.0): 160 | assert delta >= 0.0 and delta <= 255 161 | self.delta = delta 162 | 163 | def __call__(self, image, *args): 164 | if random.randint(2): 165 | image[:, :, 0] += np.uint8(random.uniform(-self.delta, self.delta)) 166 | image[:, :, 0][image[:, :, 0] > 255.0] -= 255 167 | image[:, :, 0][image[:, :, 0] < 0.0] += 255 168 | if len(args): 169 | return (image,) 170 | else: 171 | return image 172 | 173 | class RandomChannel(object): 174 | def __init__(self, delta=15.0): 175 | assert delta >= 0.0 and delta <= 255 176 | self.delta = delta 177 | 178 | def __call__(self, image, *args): 179 | channels=image.shape[-1] 180 | random_switch=np.random.randint(0,2,channels) 181 | for i in range(channels): 182 | if random_switch[i]: 183 | image[:, :, i] += np.uint8(random.uniform(-self.delta, self.delta)) 184 | image[image > 255.0] = 255 185 | image[image < 0.0] = 0 186 | if len(args): 187 | return (image,) 188 | else: 189 | return image 190 | 191 | class RandomValue(object): 192 | def __init__(self, delta=15.0): 193 | # random add or sub a random value in hsv mode 194 | assert delta>=0.0 and delta<=255.0 195 | self.delta = int(delta) 196 | def __call__(self, image, *args): 197 | if random.randint(2): 198 | image[:, :, 2] += np.uint8(random.randint(-self.delta, self.delta)) 199 | image[:, :, 2][image[:, :, 2] > 255.0] = 255 200 | image[:, :, 2][image[:, :, 2] < 0.0] = 0 201 | if len(args): 202 | return (image, *args) 203 | else: 204 | return image 205 | 206 | class RandomLightingNoise(object): 207 | def __init__(self): 208 | self.perms = ((0, 1, 2), (0, 2, 1), 209 | (1, 0, 2), (1, 2, 0), 210 | (2, 0, 1), (2, 1, 0)) 211 | 212 | def __call__(self, image, *args): 213 | if random.randint(2): 214 | swap = self.perms[random.randint(len(self.perms))] 215 | shuffle = SwapChannels(swap) # shuffle channels 216 | image = shuffle(image) 217 | if len(args): 218 | return (image, *args) 219 | else: 220 | return image 221 | 222 | 223 | class ConvertColor(object): 224 | def __init__(self, current='BGR', transform='HSV'): 225 | self.transform = transform 226 | self.current = current 227 | 228 | def __call__(self, image, *args): 229 | if self.current == 'BGR' and self.transform == 'HSV': 230 | image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) 231 | elif self.current == 'HSV' and self.transform == 'BGR': 232 | image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR) 233 | else: 234 | raise NotImplementedError 235 | if len(args): 236 | return (image, *args) 237 | else: 238 | return image 239 | 240 | 241 | class RandomContrast(object): 242 | def __init__(self, lower=0.5, upper=1.5): 243 | self.lower = lower 244 | self.upper = upper 245 | assert self.upper >= self.lower, "contrast upper must be >= lower." 246 | assert self.lower >= 0, "contrast lower must be non-negative." 247 | 248 | # expects float image 249 | def __call__(self, image,*args): 250 | if random.randint(2): 251 | alpha = random.uniform(self.lower, self.upper) 252 | image *= alpha 253 | image=image.astype(np.uint8) 254 | if len(args): 255 | return (image, ) 256 | else: 257 | return image 258 | 259 | 260 | class RandomBrightness(object): 261 | def __init__(self, delta=32): 262 | assert delta >= 0.0 263 | assert delta <= 255.0 264 | self.delta = delta 265 | 266 | def __call__(self, image, *args): 267 | if random.randint(2): 268 | delta = random.uniform(-self.delta, self.delta) 269 | image += delta 270 | if len(args): 271 | return (image, *args) 272 | else: 273 | return image 274 | 275 | class RandomNoise(object): 276 | def __init__(self,max_noise=3): 277 | self.max_noise=max_noise 278 | def __call__(self,image,*args): 279 | if np.random.randint(2): 280 | noise=np.random.randint(-self.max_noise,self.max_noise,image.shape) 281 | image=np.uint8(noise+image) 282 | if len(args): 283 | return (image,*args) 284 | else: 285 | return image 286 | 287 | class ToTensor(object): 288 | def __call__(self, cvimage, *args): 289 | image=torch.from_numpy(cvimage.astype(np.float32)).permute(2, 0, 1) 290 | if len(args): 291 | return (image, *args) 292 | else: 293 | return image 294 | 295 | 296 | class SSD_RandomSampleCrop(object): 297 | """Crop 298 | Arguments: 299 | img (Image): the image being input during training 300 | boxes (Tensor): the original bounding boxes in pt form 301 | labels (Tensor): the class labels for each bbox 302 | mode (float tuple): the min and max jaccard overlaps 303 | Return: 304 | (img, boxes, classes) 305 | img (Image): the cropped image 306 | boxes (Tensor): the adjusted bounding boxes in pt form 307 | labels (Tensor): the class labels for each bbox 308 | """ 309 | def __init__(self,sample_options=None): 310 | if sample_options is None: 311 | self.sample_options = ( 312 | # using entire original input image 313 | None, 314 | # sample a patch s.t. MIN jaccard w/ obj in .1,.3,.4,.7,.9 315 | (0.1, None), 316 | (0.3, None), 317 | (0.7, None), 318 | (0.9, None), 319 | # randomly sample a patch 320 | (None, None), 321 | ) 322 | else: 323 | self.sample_options=sample_options 324 | 325 | def __call__(self, image, boxes=None, labels=None): 326 | height, width, _ = image.shape 327 | while True: 328 | # randomly choose a mode 329 | mode = random.choice(self.sample_options) 330 | if mode is None: 331 | return image, boxes, labels 332 | 333 | min_iou, max_iou = mode 334 | if min_iou is None: 335 | min_iou = float('-inf') 336 | if max_iou is None: 337 | max_iou = float('inf') 338 | 339 | # max trails (50) 340 | for _ in range(50): 341 | current_image = image 342 | 343 | w = random.uniform(0.3 * width, width) 344 | h = random.uniform(0.3 * height, height) 345 | 346 | # aspect ratio constraint b/t .5 & 2 347 | if h / w < 0.5 or h / w > 2: 348 | continue 349 | 350 | left = random.uniform(width - w) 351 | top = random.uniform(height - h) 352 | 353 | # convert to integer rect x1,y1,x2,y2 354 | rect = np.array([int(left), int(top), int(left+w), int(top+h)]) 355 | 356 | # calculate IoU (jaccard overlap) b/t the cropped and gt boxes 357 | overlap = jaccard_numpy(boxes, rect) 358 | 359 | # is min and max overlap constraint satisfied? if not try again 360 | if overlap.min() < min_iou and max_iou < overlap.max(): 361 | continue 362 | 363 | # cut the crop from the image 364 | current_image = current_image[rect[1]:rect[3], rect[0]:rect[2], 365 | :] 366 | 367 | # keep overlap with gt box IF center in sampled patch 368 | centers = (boxes[:, :2] + boxes[:, 2:]) / 2.0 369 | 370 | # mask in all gt boxes that above and to the left of centers 371 | m1 = (rect[0] < centers[:, 0]) * (rect[1] < centers[:, 1]) 372 | 373 | # mask in all gt boxes that under and to the right of centers 374 | m2 = (rect[2] > centers[:, 0]) * (rect[3] > centers[:, 1]) 375 | 376 | # mask in that both m1 and m2 are true 377 | mask = m1 * m2 378 | 379 | # have any valid boxes? try again if not 380 | if not mask.any(): 381 | continue 382 | 383 | # take only matching gt boxes 384 | current_boxes = boxes[mask, :].copy() 385 | 386 | # take only matching gt labels 387 | current_labels = labels[mask] 388 | 389 | # should we use the box left and top corner or the crop's 390 | current_boxes[:, :2] = np.maximum(current_boxes[:, :2], 391 | rect[:2]) 392 | # adjust to crop (by substracting crop's left,top) 393 | current_boxes[:, :2] -= rect[:2] 394 | 395 | current_boxes[:, 2:] = np.minimum(current_boxes[:, 2:], 396 | rect[2:]) 397 | # adjust to crop (by substracting crop's left,top) 398 | current_boxes[:, 2:] -= rect[:2] 399 | 400 | return current_image, current_boxes, current_labels 401 | 402 | class Padding(object): 403 | def __init__(self, size=1,type='constant',constant_color=(255,255,255)): 404 | self.size = size 405 | self.type = type 406 | self.constant_color=list(constant_color) 407 | 408 | def __call__(self, image, *args): 409 | size=self.size 410 | if self.type=='constant': 411 | image = cv2.copyMakeBorder(image, size, size, size, size, cv2.BORDER_CONSTANT, value=self.constant_color) 412 | elif self.type=='reflect': 413 | image = cv2.copyMakeBorder(image, size, size, size, size, cv2.BORDER_REFLECT) 414 | elif self.type=='replicate': 415 | image = cv2.copyMakeBorder(image, size, size, size, size, cv2.BORDER_REPLICATE) 416 | 417 | if len(args): 418 | return (image, *args) 419 | else: 420 | return image 421 | 422 | class RandomCrop(object): 423 | def __init__(self,padding): 424 | self.padding=padding 425 | 426 | def __call__(self,image,*args): 427 | xi=random.randint(0,self.padding*2) 428 | yi=random.randint(0,self.padding*2) 429 | image=image[xi:-(2*self.padding-xi),yi:-(self.padding*2-yi),:] 430 | if len(args): 431 | return (image, *args) 432 | else: 433 | return image 434 | 435 | class Scale(object): 436 | def __init__(self,dim): 437 | self.dim=dim 438 | def __call__(self,image,*args): 439 | image=cv2.resize(image,self.dim) 440 | if len(args): 441 | return (image, *args) 442 | else: 443 | return image 444 | 445 | class Flip(object): 446 | def __init__(self,dim=1,percentage=0.5): 447 | """1 for horizontal flip 448 | 0 for vertical flip 449 | -1 for both flip""" 450 | self.dim=dim 451 | self.percentage=percentage 452 | 453 | def __call__(self, image, *args): 454 | if np.random.rand() Since the pytorch models we use are basically dynamic graph structures, the problem with dynamic graphs is that the graph structure cannot be determined once the forward is incomplete, and caffe is a static graph framework, which will cause the model conversion from pytorch to caffe to encounter many problems, and the pytorch version iteration is very fast,so this repo will no longer be recommended.If you want to convert pytorch to caffe,we suggest using pytorch->onnx->caffe by this repo [https://github.com/xxradon/ONNXToCaffe](https://github.com/xxradon/ONNXToCaffe "https://github.com/xxradon/ONNXToCaffe"). 3 | 4 | #### The code is mainly come from [nn_tools](https://github.com/hahnyuan/nn_tools).Thanks for hahnyuan's contribution. 5 | # Neural Network Tools: Converter and Analyser 6 | 7 | Providing a tool for neural network frameworks for pytorch and caffe. 8 | 9 | The nn_tools is released under the MIT License (refer to the LICENSE file for details). 10 | 11 | ### features 12 | 13 | 1. Converting a pytorch model to caffe model. 14 | 2. Some convenient tools of manipulate caffemodel and prototxt quickly(like get or set weights of layers). 15 | 3. Support pytorch version >= 0.2.(Have tested on 0.3,0.3.1, 0.4, 0.4.1 ,1.0, 1.2) 16 | 4. Analysing a model, get the operations number(ops) in every layers. 17 | 18 | Noting: pytorch version 1.1 is not supported now 19 | 20 | ### requirements 21 | 22 | - Python2.7 or Python3.x 23 | - Each functions in this tools requires corresponding neural network python package (pytorch and so on). 24 | 25 | # Analyser 26 | 27 | The analyser can analyse all the model layers' [input_size, output_size, multiplication ops, addition ops, 28 | comparation ops, tot ops, weight size and so on] given a input tensor size, which is convenint for model deploy analyse. 29 | 30 | ## Caffe 31 | Before you analyse your network, [Netscope](http://ethereon.github.io/netscope/#/editor) 32 | is recommended to visiualize your network. 33 | 34 | Command：`python caffe_analyser.py [-h] prototxt outdir shape` 35 | - The prototxt is the path of the prototxt file. 36 | - The outdir is path to save the csv file. 37 | - The shape is the input shape of the network(split by comma `,`), in caffe image shape should be: 38 | batch_size, channel, image_height, image_width. 39 | 40 | For example `python caffe_analyser.py resnet_18_deploy.prototxt analys_result.csv 1,3,224,224` 41 | 42 | ## Pytorch 43 | Supporting analyse the inheritors of torch.nn.Moudule class. 44 | 45 | Command：`pytorch_analyser.py [-h] [--out OUT] [--class_args ARGS] path name shape` 46 | - The path is the python file path which contaning your class. 47 | - The name is the class name or instance name in your python file. 48 | - The shape is the input shape of the network(split by comma `,`), in pytorch image shape should be: 49 | batch_size, channel, image_height, image_width. 50 | - The out (optinal) is path to save the csv file, default is '/tmp/pytorch_analyse.csv'. 51 | - The class_args (optional) is the args to init the class in python file, default is empty. 52 | 53 | For example `python pytorch_analyser.py example/resnet_pytorch_analysis_example.py resnet18 1,3,224,224` 54 | 55 | 56 | 57 | # Converter 58 | 59 | ## Pytorch to Caffe 60 | 61 | The new version of pytorch_to_caffe supporting the newest version(from 0.2.0 to 1.2.0) of pytorch. 62 | NOTICE: The transfer output will be somewhat different with the original model, caused by implementation difference. 63 | 64 | - Supporting layers types: 65 | conv2d -> Convolution, 66 | _conv_transpose2d -> Deconvolution, 67 | _linear -> InnerProduct, 68 | _split -> Slice, 69 | max_pool2d,_avg_pool2d -> Pooling, 70 | _max -> Eltwise, 71 | _cat -> Concat, 72 | dropout -> Dropout, 73 | relu -> ReLU, 74 | prelu -> PReLU, 75 | _leaky_relu -> ReLU, 76 | _tanh -> TanH, 77 | threshold(only value=0) -> Threshold,ReLU, 78 | softmax -> Softmax, 79 | batch_norm -> BatchNorm,Scale, 80 | instance_norm -> BatchNorm,Scale, 81 | _interpolate -> Upsample 82 | _hardtanh -> ReLU6 83 | _permute -> Permute 84 | _l2Norm -> Normalize 85 | 86 | 87 | - Supporting operations: torch.split, torch.max, torch.cat ,torch.sigmoid, torch.div 88 | - Supporting tensor Variable operations: var.view, + (add), += (iadd), -(sub), -=(isub) 89 | \* (mul) *= (imul), torch.Tensor.contiguous(_contiguous), torch.Tensor.pow(_pow), 90 | \* torch.Tensor.sum(_sum), torch.Tensor.sqrt(_sqrt), torch.Tensor.unsqueeze(_unsqueeze) 91 | \* torch.Tensor.expand_as(_expand_as), 92 | 93 | Need to be added for caffe in the future: 94 | - DepthwiseConv 95 | 96 | The supported above can transfer many kinds of nets, 97 | such as AlexNet(tested), VGG(tested), ResNet(fixed the bug in origin repo which mainly caused by ReLu layer function.), Inception_V3(tested). 98 | 99 | The supported layers concluded the most popular layers and operations. 100 | The other layer types will be added soon, you can ask me to add them in issues. 101 | 102 | Example: please see file `example/alexnet_pytorch_to_caffe.py`. Just Run `python3 example/alexnet_pytorch_to_caffe.py`. 103 | 104 | Attention: 105 | the main difference from convert model is the BN layer,you should pay more attention to the BN parameters like momentum=0.1, eps=1e-5. 106 | 107 | # Deploy verify(Very Important) 108 | After Converter,we should use verify_deploy.py to verify the output of pytorch model and the convertted caffe model. 109 | If you want to verify the outputs of caffe and pytorch,you should make caffe and pytorch install in the same environment,anaconda is recommended. 110 | using following script,we can install caffe-gpu(master branch). 111 | ```angular2html 112 | conda install caffe-gpu pytorch cudatoolkit=9.0 -c pytorch 113 | 114 | ``` 115 | other way,we can use docker,and in https://github.com/ufoym/deepo,for cuda9 116 | ``` 117 | docker pull ufoym/deepo:all-py36-cu90 118 | ``` 119 | for cuda10 120 | ``` 121 | docker pull ufoym/deepo:all-py36-cu100 122 | ``` 123 | 124 | please see file `example/verify_deploy.py`,it can verify the output of pytorch model and the convertted caffe model in the same input. 125 | 126 | 127 | # Some common functions 128 | 129 | ## funcs.py 130 | 131 | - **get_iou(box_a, box_b)** intersection over union of two boxes 132 | - **nms(bboxs,scores,thresh)** Non-maximum suppression 133 | - **Logger** print some str to a file and stdout with H M S 134 | 135 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/__init__.py -------------------------------------------------------------------------------- /analysis/CaffeA.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from collections import OrderedDict 3 | from .layers import * 4 | from .roi import * 5 | 6 | def profiling(net, input=None): 7 | # input is either a Blob with the shape of (batch,h,w,c) or a dict of them 8 | layers=[] 9 | if isinstance(input,dict): 10 | blob_dict = OrderedDict(input) 11 | not_ref = [input[k] for k in input] 12 | else: 13 | blob_dict = OrderedDict({'data': input}) 14 | not_ref=[input] 15 | for i, layer in enumerate(net.net.layer): 16 | out = None 17 | if len(layer.top) == 1 and len(layer.bottom) == 1: 18 | if layer.type == 'Convolution': 19 | param = layer.convolution_param 20 | out = Conv(blob_dict[layer.bottom[0]], param.kernel_size, param.num_output, param.stride, 21 | param.pad, None, layer.name, group_size=param.group) 22 | if layer.type == 'InnerProduct': 23 | param=layer.inner_product_param 24 | out= fc(blob_dict[layer.bottom[0]],param.num_output,None,layer.name) 25 | if layer.type == 'ReLU': 26 | out = Activation(blob_dict[layer.bottom[0]], 'relu', layer.name) 27 | if layer.type == 'PReLU': 28 | out = Activation(blob_dict[layer.bottom[0]], 'prelu', layer.name) 29 | if layer.type == 'Pooling': 30 | param = layer.pooling_param 31 | out = Pool(blob_dict[layer.bottom[0]], param.kernel_size, param.stride, 32 | param.pad, layer.name,param.pool,ceil=True) 33 | if layer.type == 'Normalize': 34 | out = Norm(blob_dict[layer.bottom[0]], 'norm', layer.name) 35 | if layer.type == 'BatchNorm': 36 | out= Norm(blob_dict[layer.bottom[0]],'batch_norm',layer.name) 37 | if layer.type== 'LRN': 38 | out= Norm(blob_dict[layer.bottom[0]],'lrn',layer.name) 39 | if layer.type == 'Permute': 40 | shape=[blob_dict[layer.bottom[0]][dim-1] for dim in layer.permute_param.order[1:]] 41 | out = Permute(blob_dict[layer.bottom[0]],shape,layer.name) 42 | if layer.type == 'Flatten': 43 | out = Flatten(blob_dict[layer.bottom[0]], layer.name) 44 | if layer.type == 'Scale': 45 | out =Scale (blob_dict[layer.bottom[0]], name = layer.name) 46 | if layer.type == 'Softmax': 47 | out =Softmax (blob_dict[layer.bottom[0]], name = layer.name) 48 | if layer.type == 'Dropout': 49 | out =Dropout (blob_dict[layer.bottom[0]], name = layer.name) 50 | if layer.type == 'Reshape': 51 | out =Reshape (blob_dict[layer.bottom[0]],shape=layer.reshape_param.shape.dim, name = layer.name) 52 | if out: 53 | try: 54 | not_ref.remove(blob_dict[layer.bottom[0]]) 55 | except: 56 | pass 57 | blob_dict[layer.top[0]] = out() 58 | not_ref.append(blob_dict[layer.top[0]]) 59 | layers.append(out) 60 | else: 61 | assert 'layer type: %s cannot be P' % (layer.type) 62 | elif len(layer.bottom)>1: 63 | # for multi input layer 64 | if layer.type=='Eltwise': 65 | param=layer.eltwise_param 66 | out = Eltwise([blob_dict[bottom] for bottom in layer.bottom], 67 | type=param.EltwiseOp.Name(param.operation),name=layer.name) 68 | if layer.type=='PSROIPooling': 69 | param=layer.psroi_pooling_param 70 | out = PSROIPool(blob_dict[layer.bottom[0]],blob_dict[layer.bottom[1]], 71 | param.output_dim,param.group_size) 72 | if layer.type=='ROIPooling': 73 | param=layer.roi_pooling_param 74 | out = ROIPool(blob_dict[layer.bottom[0]],blob_dict[layer.bottom[1]], 75 | param.pooled_w,param.pooled_h,layer.name) 76 | if layer.type == "Concat": 77 | param = layer.concat_param 78 | out = Concat([blob_dict[bottom] for bottom in layer.bottom],param.axis,layer.name) 79 | if out: 80 | for bottom in layer.bottom: 81 | try: 82 | not_ref.remove(blob_dict[bottom]) 83 | except: 84 | pass 85 | blob_dict[layer.top[0]] = out() 86 | not_ref.append(blob_dict[layer.top[0]]) 87 | layers.append(out) 88 | else: 89 | assert 'layer type: %s cannot be P' % (layer.type) 90 | elif len(layer.top)>1: 91 | if layer.type == 'Slice': 92 | param=layer.slice_param 93 | out =Slice (blob_dict[layer.bottom[0]], name = layer.name,slice_point=param.slice_point,axis=param.axis) 94 | if out: 95 | try: 96 | not_ref.remove(blob_dict[layer.bottom[0]]) 97 | except: 98 | pass 99 | for o,top in zip(out(),layer.top): 100 | blob_dict[top] = o 101 | not_ref.append(blob_dict[top]) 102 | layers.append(out) 103 | return blob_dict,layers -------------------------------------------------------------------------------- /analysis/PytorchA.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | import torch 3 | import torch.nn as nn 4 | import numpy as np 5 | from torch.autograd import Variable 6 | from collections import OrderedDict 7 | from .layers import * 8 | from . import save_csv 9 | 10 | tracked_layers=[] 11 | blob_dict=[] 12 | layer_name_dict={} 13 | 14 | def _analyse(module,raw_input): 15 | input=[] 16 | for i in raw_input: 17 | s = i.size() 18 | # if len(s)==4: 19 | # input.append(Blob([s[0],s[2],s[3],s[1]])) 20 | # else: 21 | input.append(Blob(s)) 22 | out=None 23 | name=layer_name_dict[module] 24 | if isinstance(module,nn.Conv2d): 25 | out=Conv(input[0],module.kernel_size,module.out_channels, 26 | module.stride,module.padding,group_size=module.groups,name=name) 27 | elif isinstance(module,nn.ConvTranspose2d): 28 | out=Conv(input[0],module.kernel_size,module.out_channels, 29 | module.stride,module.padding,group_size=module.groups,name=name,transpose=True) 30 | elif isinstance(module,nn.BatchNorm2d): 31 | out=Norm(input[0],'batch_norm',name=name) 32 | elif isinstance(module,nn.Linear): 33 | out=fc(input[0],module.out_features,name=name) 34 | elif isinstance(module,nn.MaxPool2d): 35 | out = pool(input[0], module.kernel_size,module.stride,module.padding, 36 | name=name,pool_type='max') 37 | elif isinstance(module,nn.AvgPool2d): 38 | out = pool(input[0], module.kernel_size,module.stride,module.padding, 39 | name=name,pool_type='avg') 40 | elif isinstance(module,nn.ReLU): 41 | out = Activation(input[0],'relu',name=name) 42 | elif isinstance(module,nn.Conv3d): 43 | out=Conv(input[0],module.kernel_size,module.out_channels, 44 | module.stride,module.padding,group_size=module.groups,name=name) 45 | 46 | if out: 47 | tracked_layers.append(out) 48 | else: 49 | print('WARNING: skip Module {}' .format(module)) 50 | 51 | def module_hook(module, input, output): 52 | # print('module hook') 53 | # print module 54 | # for i in input: 55 | # print ('input',i.size()) 56 | # for i in output: 57 | # print('out', i.size()) 58 | _analyse(module,input) 59 | 60 | def register(module): 61 | module.register_forward_hook(module_hook) 62 | 63 | def analyse(net, inputs): 64 | """ 65 | analyse the network given input 66 | :param net: torch.nn.Module 67 | :param inputs: torch.Variable, torch.Tensor or list of them 68 | :return: blob_dict, tracked_layers 69 | """ 70 | del tracked_layers[:] 71 | del blob_dict[:] 72 | if inputs is not list: 73 | raw_inputs=[inputs] 74 | _inputs=[] 75 | for name,layer in net.named_modules(): 76 | layer_name_dict[layer]=name 77 | for i in raw_inputs: 78 | if isinstance(i,Variable): 79 | _inputs.append(i) 80 | elif isinstance(i,torch.Tensor): 81 | _inputs.append(Variable(i)) 82 | elif isinstance(i,np.ndarray): 83 | _inputs.append(Variable(torch.Tensor(i))) 84 | else: 85 | raise NotImplementedError("Not Support the input type {}".format(type(i))) 86 | net.apply(register) 87 | net.forward(*_inputs) 88 | return blob_dict,tracked_layers 89 | 90 | def profilling(net,input): 91 | """ Old API of analyse """ 92 | return analyse(net,input) -------------------------------------------------------------------------------- /analysis/README.md: -------------------------------------------------------------------------------- 1 | # Blob 2 | the input Blob must have the first dimension for batch_size 3 | for image (batch_size,w,h,c) -------------------------------------------------------------------------------- /analysis/__init__.py: -------------------------------------------------------------------------------- 1 | from .layers import * 2 | from .roi import * 3 | from .utils import * 4 | from .blob import * -------------------------------------------------------------------------------- /analysis/__pycache__/CaffeA.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/analysis/__pycache__/CaffeA.cpython-36.pyc -------------------------------------------------------------------------------- /analysis/__pycache__/PytorchA.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/analysis/__pycache__/PytorchA.cpython-36.pyc -------------------------------------------------------------------------------- /analysis/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/analysis/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /analysis/__pycache__/blob.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/analysis/__pycache__/blob.cpython-36.pyc -------------------------------------------------------------------------------- /analysis/__pycache__/layers.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/analysis/__pycache__/layers.cpython-36.pyc -------------------------------------------------------------------------------- /analysis/__pycache__/roi.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/analysis/__pycache__/roi.cpython-36.pyc -------------------------------------------------------------------------------- /analysis/__pycache__/utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/analysis/__pycache__/utils.cpython-36.pyc -------------------------------------------------------------------------------- /analysis/blob.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | # the blob with shape of (c,h,w) or (batch,c,h,w) for image 4 | # the blob with shape of (batch,c,t,h,w) for video 5 | class Blob(): 6 | def __init__(self,shape,father=None): 7 | shape=[int(i) for i in shape] 8 | self._data=None 9 | self.shape=[int(i) for i in list(shape)] 10 | self.father=type(father)==list and father or [father] 11 | 12 | @property 13 | def data(self): 14 | raise NotImplementedError('Blob.data is removed from this version of nn_tools, you should use .shape') 15 | 16 | @property 17 | def size(self): 18 | return np.prod(self.shape) 19 | 20 | @property 21 | def w(self): 22 | if self.dim in [3,4,5]: 23 | return self.shape[-1] 24 | else: 25 | raise NotImplementedError('Blob attribute w is only supported for 2D or 3D feature map') 26 | 27 | @property 28 | def h(self): 29 | if self.dim in [3,4,5]: 30 | return self.shape[-2] 31 | else: 32 | raise NotImplementedError('Blob attribute h is only supported for 2D or 3D feature map') 33 | 34 | @property 35 | def t(self): 36 | # time attribute 37 | if self.dim in [5]: 38 | # 3D feature map 39 | return self.shape[3] 40 | else: 41 | raise NotImplementedError('Blob attribute t is only supported for 3D feature map') 42 | 43 | @property 44 | def c(self): 45 | if self.dim in [3,4]: 46 | # 2D feature map 47 | return self.shape[-3] 48 | elif self.dim in [5]: 49 | # 3D feature map 50 | return self.shape[1] 51 | else: 52 | raise NotImplementedError('Blob attribute h is only supported for 2D or 3D feature map') 53 | 54 | @property 55 | def batch_size(self): 56 | return self.shape[0] 57 | 58 | @property 59 | def dim(self): 60 | return len(self.shape) 61 | 62 | def new(self,father): 63 | return Blob(self.shape,father) 64 | 65 | def __getitem__(self, key): 66 | return self.shape[key] 67 | 68 | def __str__(self): 69 | return str(self.shape) 70 | 71 | def flaten(self): 72 | return Blob([np.prod(self.shape)]) -------------------------------------------------------------------------------- /analysis/layers.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from .blob import Blob 3 | 4 | 5 | box=[] 6 | 7 | class Base(object): 8 | def __init__(self,input,name=''): 9 | def transfer_input(_input): 10 | if isinstance(_input,Base): 11 | _input=_input() 12 | assert isinstance(input,Blob),'The input of layer %s is not Blob, please use nn_tools.P.blob.Blob as input'%name 13 | return _input 14 | if type(input)==list: 15 | # if multi input 16 | self.input=[transfer_input(i) for i in input] 17 | self.input_size = np.sum([np.prod(i.shape) for i in self.input]) 18 | self.muti_input=True 19 | else: 20 | self.input = transfer_input(input) 21 | self.input_size = np.prod(self.input.shape) 22 | self.muti_input = False 23 | self.name=name 24 | self.weight_size=0 25 | self.activation_size=None 26 | self.dot=0 27 | self.add=0 28 | self.pow=0 29 | self.compare=0 30 | self.ops=0 31 | self.out=None 32 | self.layer_info=None 33 | box.append(self) 34 | 35 | def __call__(self, *args, **kwargs): 36 | return self.out 37 | def __setattr__(self, key, value): 38 | if key=='out' and value!=None: 39 | if type(value) is list: 40 | self.activation_size=0 41 | for i in value: 42 | self.activation_size+=np.prod(i.shape) 43 | else: 44 | self.activation_size=np.prod(value.shape) 45 | return object.__setattr__(self, key,value) 46 | def __getattribute__(self, item): 47 | if item=='ops': 48 | try: 49 | self.ops=self.pow+self.add+self.dot+self.compare 50 | except: 51 | print("CRITICAL WARNING: Layer {} ops cannot be calculated, set to 0.".format(self.name)) 52 | self.ops=0 53 | return object.__getattribute__(self,item) 54 | 55 | class Norm(Base): 56 | valid_tuple=('norm','batch_norm','lrn') 57 | def __init__(self,input,type,name=None): 58 | if type not in Norm.valid_tuple: 59 | raise NameError('the norm type:' + type + ' is not supported. ' \ 60 | 'the valid type is: ' + str(Activation.valid_tuple)) 61 | if name == None: name = type 62 | Base.__init__(self, input, name=name) 63 | getattr(self, type)() 64 | self.out = self.input.new(self) 65 | 66 | def norm(self): 67 | self.dot = self.input_size 68 | self.add = self.input_size 69 | 70 | def batch_norm(self): 71 | self.dot = self.input_size 72 | self.add = self.input_size 73 | 74 | def lrn(self): 75 | self.dot = self.input_size 76 | self.add = self.input_size 77 | 78 | class Activation(Base): 79 | #valid tuple lists the valid activation function type 80 | valid_tuple=('relu','tanh','prelu') 81 | def __init__(self,input,type,name=None): 82 | if type not in Activation.valid_tuple: 83 | raise NameError('the activation type:'+type+' is not supported. ' \ 84 | 'the valid type is: '+str(Activation.valid_tuple)) 85 | if name==None:name=type 86 | Base.__init__(self,input,name=name) 87 | getattr(self,type)() 88 | self.out=self.input.new(self) 89 | 90 | def relu(self): 91 | self.compare=self.input_size 92 | 93 | def sigmoid(self): 94 | self.add=self.dot=self.pow=self.input_size 95 | 96 | def tanh(self): 97 | self.dot=self.input_size 98 | self.add=self.pow=self.input_size*2 99 | 100 | def prelu(self): 101 | self.compare=self.input_size 102 | self.dot=self.input_size 103 | 104 | 105 | class Sliding(Base): 106 | def __init__(self,input,kernel_size,num_out,stride=1,pad=0,name='sliding',ceil=False,transpose=False): 107 | # input is the instance of blob.Blob with shape (c,h,w) or (batch,c,h,w) 108 | super(Sliding,self).__init__(input,name=name) 109 | if self.input.dim==4: 110 | conv_dims=2 111 | elif self.input.dim==5: 112 | conv_dims=3 113 | self.input_t=self.input.t 114 | else: 115 | raise ValueError('Sliding must have a input with 2D Map(batch,c,w,h) or 3D Map(batch,c,d,w,h)') 116 | self.input_w = self.input.w 117 | self.input_h = self.input.h 118 | self.batch_size = self.input.batch_size 119 | self.in_channel = self.input.c 120 | 121 | if type(kernel_size) == int: 122 | self.kernel_size = [kernel_size] * conv_dims 123 | else: 124 | assert len(kernel_size)==conv_dims 125 | self.kernel_size = [i for i in kernel_size] 126 | if type(stride) == int: 127 | self.stride = [stride] * conv_dims 128 | else: 129 | self.stride = [i for i in stride] 130 | if len(self.stride) == 1: 131 | self.stride = [self.stride[0]] * conv_dims 132 | elif len(self.stride) == 0: 133 | self.stride = [1] * conv_dims 134 | if type(pad) == int: 135 | self.pad = [pad] * conv_dims 136 | else: 137 | self.pad = [i for i in pad] 138 | if len(self.pad) == 1: 139 | self.pad *= conv_dims 140 | elif len(self.pad) == 0: 141 | self.pad = [0] * conv_dims 142 | self.num_out = num_out 143 | self.layer_info ='kernel=%s,stride=%s,pad=%s'%('x'.join([str(_) for _ in self.kernel_size]), 144 | 'x'.join([str(_) for _ in self.stride]), 145 | 'x'.join([str(_) for _ in self.pad])) 146 | if transpose: 147 | self.layer_info += ',transpose' 148 | # calc out 149 | 150 | outs=[] 151 | 152 | for i in range(self.input.dim-2): 153 | if not transpose: 154 | if not ceil: 155 | outs.append(np.floor(float(self.input[2+i] + self.pad[i] * 2 - self.kernel_size[i]) / self.stride[i]) + 1) 156 | else: 157 | outs.append(np.ceil(float(self.input[2+i] + self.pad[i] * 2 - self.kernel_size[i]) / self.stride[i]) + 1) 158 | else: 159 | # transpose 160 | outs.append((self.input[2+i] - 1) * self.stride[i] - 2 * self.pad[i] + self.kernel_size[i]) 161 | # if not ceil: 162 | # out_h = np.floor(float(self.input_w + self.pad[0] * 2 - self.kernel_size[0]) / self.stride[0]) + 1 163 | # out_w = np.floor(float(self.input_h + self.pad[1] * 2 - self.kernel_size[1]) / self.stride[1]) + 1 164 | # out_t = np.floor(float(self.input_t + self.pad[2] * 2 - self.kernel_size[2]) / self.stride[2]) + 1 165 | # else: 166 | # out_w = np.ceil(float(self.input_w + self.pad[0] * 2 - self.kernel_size[0]) / self.stride[0]) + 1 167 | # out_h = np.ceil(float(self.input_h + self.pad[1] * 2 - self.kernel_size[1]) / self.stride[1]) + 1 168 | # out_t = np.ceil(float(self.input_h + self.pad[1] * 2 - self.kernel_size[1]) / self.stride[1]) + 1 169 | # else: 170 | # # transpose 171 | # out_w = (self.input_w - 1) * self.stride[0] - 2 * self.pad[0] + self.kernel_size[0] 172 | # out_h = (self.input_h - 1) * self.stride[1] - 2 * self.pad[1] + self.kernel_size[1] 173 | 174 | self.out = Blob([self.batch_size, num_out, *outs], self) 175 | 176 | class Conv(Sliding): 177 | def __init__(self,input,kernel_size,num_out,stride=1,pad=0, 178 | activation='relu',name='conv',ceil=False,group_size=1,transpose=False): 179 | if isinstance(input,Base): 180 | input=input() 181 | Sliding.__init__(self,input,kernel_size,num_out,stride,pad,name=name,ceil=ceil,transpose=transpose) 182 | self.layer_info+=',num_out=%d'%(num_out) 183 | self.dot = np.prod(self.out.shape) * np.prod(self.kernel_size) * self.in_channel 184 | self.weight_size = np.prod(self.kernel_size) * num_out * self.in_channel 185 | if group_size!=1: 186 | self.layer_info += ',group_size=%d' % (group_size) 187 | self.dot /= group_size 188 | self.weight_size /= group_size 189 | self.add = self.dot 190 | if activation: 191 | Activation(self.out,activation) 192 | 193 | class Pool(Sliding): 194 | def __init__(self,input,kernel_size,stride=1,pad=0,name='pool',pool_type='max',ceil=False): 195 | # pool_type: 0 is max, 1 is avg/ave in Caffe 196 | if isinstance(input,Base): 197 | input=input() 198 | Sliding.__init__(self,input,kernel_size,input.c,stride,pad,name=name,ceil=ceil) 199 | self.pool_type=pool_type 200 | self.layer_info+=',type=%s'%(pool_type) 201 | if pool_type in ['max',0]: 202 | self.compare= np.prod(self.out.shape) * (np.prod(self.kernel_size) - 1) 203 | elif pool_type in ['avg','ave',1]: 204 | self.add = np.prod(self.input.shape) 205 | self.dot = np.prod(self.out.shape) 206 | else: 207 | print("WARNING, NOT IMPLEMENT POOL TYPE %s PROFILING at %s, CONTINUE"%(pool_type,name)) 208 | pool=Pool 209 | 210 | class InnerProduct(Base): 211 | def __init__(self,input,num_out,activation='relu',name='innerproduct'): 212 | if isinstance(input,Base): 213 | input=input() 214 | Base.__init__(self,input,name=name) 215 | self.left_dim=np.prod(input.shape[1:]) 216 | self.num_out=num_out 217 | self.dot=self.num_out*self.input_size 218 | self.add=self.num_out*self.input_size 219 | self.out=Blob([input[0],self.num_out],self) 220 | self.weight_size = self.num_out * self.left_dim 221 | if activation: 222 | Activation(self.out,activation) 223 | Fc=InnerProduct 224 | fc=InnerProduct 225 | 226 | class Permute(Base): 227 | def __init__(self, input,dims, name='permute'): 228 | super(Permute,self).__init__(input,name) 229 | self.out = Blob(dims,self) 230 | 231 | class Flatten(Base): 232 | def __init__(self,input, name='permute'): 233 | super(Flatten, self).__init__(input, name) 234 | dim=[np.prod(input.shape)] 235 | self.out = Blob(dim, self) 236 | 237 | class Eltwise(Base): 238 | def __init__(self,inputs,type='sum',name='eltwise'): 239 | super(Eltwise,self).__init__(inputs,name,) 240 | self.out=inputs[0].new(self) 241 | if type in ['sum','SUM']: 242 | self.add=np.prod(self.out.shape) 243 | elif type in ['product','PROD']: 244 | self.dot=np.prod(self.out.shape) 245 | elif type in ['max','MAX']: 246 | self.compare=np.prod(self.out.shape) 247 | else: 248 | raise AttributeError('the Eltwise layer type must be sum, max or product') 249 | 250 | class Slice(Base): 251 | def __init__(self,input,slice_point,axis,name='slice'): 252 | super(Slice,self).__init__(input,name,) 253 | self.out=[] 254 | last=0 255 | for p in slice_point: 256 | print(p,list(input.shape)) 257 | shape1=list(input.shape) 258 | shape1[axis] = p-last 259 | last=p 260 | self.out+=[Blob(shape1)] 261 | shape1 = list(input.shape) 262 | print(last,shape1,input.shape[axis]) 263 | shape1[axis] = input.shape[axis] - last 264 | self.out += [Blob(shape1)] 265 | 266 | class Reshape(Base): 267 | def __init__(self,input,shape,name='reshape'): 268 | super(Reshape,self).__init__(input,name) 269 | shape=list(shape) 270 | for i in range(len(shape)): 271 | if shape[i]==0: 272 | shape[i]=input.shape[i] 273 | self.out=Blob(shape) 274 | 275 | 276 | class Concat(Base): 277 | def __init__(self,inputs,axis,name='concat'): 278 | super(Concat,self).__init__(inputs,name,) 279 | outc=0 280 | for input in inputs: 281 | outc+=input[axis] 282 | self.out=Blob(inputs[0].shape,self) 283 | self.out.shape[axis]=outc 284 | 285 | class Scale(Base): 286 | def __init__(self, input, factor=None, name='scale'): 287 | super(Scale, self).__init__(input, name, ) 288 | self.out = input.new(self) 289 | 290 | self.dot=self.input_size 291 | # TODO scale analysis 292 | 293 | class Softmax(Base): 294 | def __init__(self, input, factor=None, name='softmax'): 295 | super(Softmax, self).__init__(input, name, ) 296 | self.out = input.new(self) 297 | self.power=self.input_size 298 | self.add=self.input_size 299 | self.dot=self.input_size 300 | self.layer_info="softmax" 301 | 302 | class Dropout(Base): 303 | def __init__(self,input,name='dropout'): 304 | if isinstance(input,Base): 305 | input=input() 306 | Base.__init__(self,input,name=name) 307 | self.out = input.new(self) -------------------------------------------------------------------------------- /analysis/roi.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from .blob import Blob 3 | from .layers import Base 4 | 5 | 6 | class Flatten(Base): 7 | def __init__(self,input, name='permute'): 8 | super(Flatten, self).__init__(input, name) 9 | dim=[np.prod(input.data.shape)] 10 | self.out = Blob(dim, self) 11 | 12 | class PSROIPool(Base): 13 | def __init__(self,input,rois,output_dim,group_size,name='psroipool'): 14 | super(PSROIPool,self).__init__([input,rois],name) 15 | self.rois=rois 16 | dim=[rois.shape[0],output_dim,group_size,group_size] 17 | self.out=Blob(dim,self) 18 | self.layer_info='output_dim:%d,group_size:%d'%(output_dim,group_size) 19 | 20 | # TODO PSROIPOOL ANALYSIS 21 | 22 | class ROIPool(Base): 23 | def __init__(self,input,rois,pooled_w,pooled_h,name='roipool'): 24 | super(ROIPool,self).__init__([input,rois],name) 25 | self.rois = rois 26 | dim=[rois.shape[0],pooled_w,pooled_h,input[3]] 27 | self.out = Blob(dim, self) 28 | self.layer_info = 'roi pooled:%dx%d' % (pooled_w, pooled_h) 29 | 30 | # TODO PSROIPOOL ANALYSIS -------------------------------------------------------------------------------- /analysis/utils.py: -------------------------------------------------------------------------------- 1 | import csv,pprint 2 | from .layers import Base 3 | 4 | def get_human_readable(num): 5 | units=['','K','M','G','T','P'] 6 | idx=0 7 | while .001*num>1: 8 | num=.001*num 9 | idx+=1 10 | if idx>=len(units): 11 | return '%.3e'%num 12 | return '%.3f'%num+units[idx] 13 | 14 | def save_csv(layers,csv_save_path='/tmp/analyse.csv', 15 | save_items=('name', 'layer_info', 'input', 'out', 'dot', 'add', 'compare','ops', 'weight_size','activation_size'), 16 | print_detail=True,human_readable=True): 17 | # layers = get_layer_blox_from_blobs(blobs) 18 | print_list = [] 19 | sum=[0]*len(save_items) 20 | for layer in layers: 21 | print_line=[] 22 | for idx,param in enumerate(save_items): 23 | item=getattr(layer, param) 24 | if type(item)==list: 25 | s='' 26 | for i in item: 27 | s+=' '+str(i) 28 | else: 29 | s=str(item) 30 | try: 31 | num=int(item) 32 | sum[idx]+=num 33 | except:pass 34 | print_line.append(s) 35 | print_list.append(print_line) 36 | if print_detail: 37 | sum[0] = 'SUM' 38 | print_list.append(sum) 39 | print_table(print_list,save_items) 40 | else: 41 | print_list=[] 42 | for idx,item in enumerate(sum): 43 | if item>0: 44 | if human_readable: 45 | print_list.append('%s:%s' % (save_items[idx], get_human_readable(item))) 46 | else: 47 | print_list.append('%s:%.3e'%(save_items[idx],item)) 48 | print(print_list) 49 | if csv_save_path!=None: 50 | with open(csv_save_path,'w') as file: 51 | writer=csv.writer(file) 52 | writer.writerow(save_items) 53 | for layer in print_list: 54 | writer.writerow(layer) 55 | print('saved at {}!'.format(csv_save_path)) 56 | 57 | def get_layer_blox_from_blobs(blobs): 58 | layers=[] 59 | def creator_search(blob): 60 | for father in blob.father: 61 | if isinstance(father,Base) and father not in layers: 62 | layers.append(father) 63 | if father.muti_input==True: 64 | for input in father.input: 65 | creator_search(input) 66 | else: 67 | creator_search(father.input) 68 | for blob in blobs: 69 | creator_search(blob) 70 | return layers 71 | 72 | def print_table(datas,names): 73 | 74 | types=[] 75 | for i in datas[0]: 76 | try: 77 | i=int(float(i)) 78 | types.append('I') 79 | except: 80 | types.append('S') 81 | for l in datas: 82 | s='' 83 | for i,t in zip(l,types): 84 | if t=='I': 85 | 86 | i=int(float(i)) 87 | s+=('%.1E'%i).center(10) 88 | else: 89 | i=str(i) 90 | if len(i)>20: 91 | i=i[:17]+'...' 92 | s+=i.center(20) 93 | s+='|' 94 | print(s) 95 | s = '' 96 | for i,t in zip(names,types): 97 | 98 | if t == 'I': 99 | s += i.center(10) 100 | else: 101 | if len(i) > 20: 102 | i = i[:17] + '...' 103 | s += i.center(20) 104 | s += '|' 105 | print(s) 106 | 107 | def print_by_blob(blobs,print_items=('name', 'layer_info', 'input', 'out', 'dot', 'add', 'compare','ops', 'weight_size','activation_size')): 108 | layers=get_layer_blox_from_blobs(blobs) 109 | print_list = [] 110 | for layer in layers: 111 | print_list.append([str(getattr(layer, param)) for param in print_items]) 112 | pprint.pprint(print_list, depth=3, width=200) 113 | return print_list -------------------------------------------------------------------------------- /caffe_analyser.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | from __future__ import absolute_import 3 | import argparse 4 | from analysis.CaffeA import * 5 | from Caffe import caffe_net 6 | from analysis.utils import save_csv 7 | 8 | """ 9 | Before you analyse your network, [Netscope](http://ethereon.github.io/netscope/#/editor) 10 | is recommended to visiualize your network. 11 | 12 | Command：`python caffe_analyser.py [-h] prototxt outdir shape` 13 | - The prototxt is the path of the prototxt file. 14 | - The outdir is path to save the csv file. 15 | - The shape is the input shape of the network(split by comma `,`), image shape should be: channel, batch_size, image_height, image_width. 16 | 17 | For example `python caffe_analyser.py resnet_18_deploy.prototxt analys_result.csv 1,3,224,224` 18 | """ 19 | 20 | if __name__=="__main__": 21 | parser=argparse.ArgumentParser() 22 | parser.add_argument('prototxt',help='path of the prototxt file',type=str) 23 | parser.add_argument('outdir',help='path to save the csv file',type=str) 24 | parser.add_argument('shape',help='input shape of the network(split by comma `,`), image shape should be: batch,c,h,w',type=str) 25 | args=parser.parse_args() 26 | shape=[int(i) for i in args.shape.split(',')] 27 | net=caffe_net.Prototxt(args.prototxt) 28 | blob_dict, layers=profiling(net, Blob(shape)) 29 | save_csv(layers,args.outdir) -------------------------------------------------------------------------------- /example/MobileNetV2.py: -------------------------------------------------------------------------------- 1 | 2 | #From https://github.com/tonylins/pytorch-mobilenet-v2 3 | #model https://drive.google.com/open?id=1jlto6HRVD3ipNkAl1lNhDbkBp7HylaqR 4 | import torch.nn as nn 5 | import math 6 | 7 | 8 | def conv_bn(inp, oup, stride): 9 | return nn.Sequential( 10 | nn.Conv2d(inp, oup, 3, stride, 1, bias=False), 11 | nn.BatchNorm2d(oup), 12 | nn.ReLU6(inplace=True) 13 | ) 14 | 15 | 16 | def conv_1x1_bn(inp, oup): 17 | return nn.Sequential( 18 | nn.Conv2d(inp, oup, 1, 1, 0, bias=False), 19 | nn.BatchNorm2d(oup), 20 | nn.ReLU6(inplace=True) 21 | ) 22 | 23 | 24 | class InvertedResidual(nn.Module): 25 | def __init__(self, inp, oup, stride, expand_ratio): 26 | super(InvertedResidual, self).__init__() 27 | self.stride = stride 28 | assert stride in [1, 2] 29 | 30 | hidden_dim = round(inp * expand_ratio) 31 | self.use_res_connect = self.stride == 1 and inp == oup 32 | 33 | if expand_ratio == 1: 34 | self.conv = nn.Sequential( 35 | # dw 36 | nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False), 37 | nn.BatchNorm2d(hidden_dim), 38 | nn.ReLU6(inplace=True), 39 | # pw-linear 40 | nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), 41 | nn.BatchNorm2d(oup), 42 | ) 43 | else: 44 | self.conv = nn.Sequential( 45 | # pw 46 | nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False), 47 | nn.BatchNorm2d(hidden_dim), 48 | nn.ReLU6(inplace=True), 49 | # dw 50 | nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False), 51 | nn.BatchNorm2d(hidden_dim), 52 | nn.ReLU6(inplace=True), 53 | # pw-linear 54 | nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), 55 | nn.BatchNorm2d(oup), 56 | ) 57 | 58 | def forward(self, x): 59 | if self.use_res_connect: 60 | return x + self.conv(x) 61 | else: 62 | return self.conv(x) 63 | 64 | 65 | class MobileNetV2(nn.Module): 66 | def __init__(self, n_class=1000, input_size=224, width_mult=1.): 67 | super(MobileNetV2, self).__init__() 68 | block = InvertedResidual 69 | input_channel = 32 70 | last_channel = 1280 71 | interverted_residual_setting = [ 72 | # t, c, n, s 73 | [1, 16, 1, 1], 74 | [6, 24, 2, 2], 75 | [6, 32, 3, 2], 76 | [6, 64, 4, 2], 77 | [6, 96, 3, 1], 78 | [6, 160, 3, 2], 79 | [6, 320, 1, 1], 80 | ] 81 | 82 | # building first layer 83 | assert input_size % 32 == 0 84 | input_channel = int(input_channel * width_mult) 85 | self.last_channel = int(last_channel * width_mult) if width_mult > 1.0 else last_channel 86 | self.features = [conv_bn(3, input_channel, 2)] 87 | # building inverted residual blocks 88 | for t, c, n, s in interverted_residual_setting: 89 | output_channel = int(c * width_mult) 90 | for i in range(n): 91 | if i == 0: 92 | self.features.append(block(input_channel, output_channel, s, expand_ratio=t)) 93 | else: 94 | self.features.append(block(input_channel, output_channel, 1, expand_ratio=t)) 95 | input_channel = output_channel 96 | # building last several layers 97 | self.features.append(conv_1x1_bn(input_channel, self.last_channel)) 98 | # make it nn.Sequential 99 | self.features = nn.Sequential(*self.features) 100 | 101 | # building classifier 102 | self.classifier = nn.Sequential( 103 | nn.Dropout(0.2), 104 | nn.Linear(self.last_channel, n_class), 105 | ) 106 | 107 | self._initialize_weights() 108 | 109 | def forward(self, x): 110 | x = self.features(x) 111 | x = x.mean(3).mean(2) 112 | x = self.classifier(x) 113 | return x 114 | 115 | def _initialize_weights(self): 116 | for m in self.modules(): 117 | if isinstance(m, nn.Conv2d): 118 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 119 | m.weight.data.normal_(0, math.sqrt(2. / n)) 120 | if m.bias is not None: 121 | m.bias.data.zero_() 122 | elif isinstance(m, nn.BatchNorm2d): 123 | m.weight.data.fill_(1) 124 | m.bias.data.zero_() 125 | elif isinstance(m, nn.Linear): 126 | n = m.weight.size(1) 127 | m.weight.data.normal_(0, 0.01) 128 | m.bias.data.zero_() 129 | -------------------------------------------------------------------------------- /example/Resnet-annalyze.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | 3 | from __future__ import absolute_import 4 | import sys 5 | sys.path.insert(0,'.') 6 | import torch 7 | import argparse 8 | import sys,os 9 | from analysis.PytorchA import analyse 10 | from analysis.utils import save_csv 11 | from torch.autograd import Variable 12 | import torch.nn as nn 13 | import sys 14 | sys.path.insert(0,'.') 15 | from model import mgn 16 | from torchvision.models import resnet 17 | """ 18 | Supporting analyse the inheritors of torch.nn.Moudule class. 19 | 20 | Command：`pytorch_analyser.py [-h] [--out OUT] [--class_args ARGS] path class_name shape` 21 | - The path is the python file path which contaning your class. 22 | - The class_name is the class name in your python file. 23 | - The shape is the input shape of the network(split by comma `,`), in pytorch image shape should be: batch_size, channel, image_height, image_width. 24 | - The out (optinal) is path to save the csv file, default is '/tmp/pytorch_analyse.csv'. 25 | - The class_args (optional) is the args to init the class in python file, default is empty. 26 | 27 | For example `python pytorch_analyser.py tmp/pytorch_analysis_test.py ResNet218 1,3,224,224` 28 | """ 29 | 30 | 31 | if __name__=="__main__": 32 | parser=argparse.ArgumentParser() 33 | parser.add_argument('--path',help='python file location',default = "MGN_analysis_example.py",type=str) 34 | parser.add_argument('--name',help='the class name or instance name in your python file',default = "mgn",type=str) 35 | parser.add_argument('--shape',help='input shape of the network(split by comma `,`), image shape should be: batch,c,h,w',type=str) 36 | parser.add_argument('--out',help='path to save the csv file',default='pytorch_analyse.csv',type=str) 37 | parser.add_argument('--class_args',help='args to init the class in python file',default='net',type=str) 38 | 39 | args=parser.parse_args() 40 | name='resnet18' 41 | net=resnet.resnet18() 42 | net.eval() 43 | #For pytorch 0.3 44 | #input=Variable(torch.ones([1,3,224,224])) 45 | 46 | x = torch.rand( [1,3,224, 224]) 47 | blob_dict, layers = analyse(net, x) 48 | save_csv(layers, args.out) 49 | 50 | 51 | -------------------------------------------------------------------------------- /example/__pycache__/MGN_analysis_example.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/example/__pycache__/MGN_analysis_example.cpython-36.pyc -------------------------------------------------------------------------------- /example/__pycache__/option.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/example/__pycache__/option.cpython-36.pyc -------------------------------------------------------------------------------- /example/alexnet_pytorch_to_caffe.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.path.insert(0,'.') 3 | import torch 4 | from torch.autograd import Variable 5 | from torchvision.models.alexnet import alexnet 6 | import pytorch_to_caffe 7 | 8 | if __name__=='__main__': 9 | name='alexnet' 10 | net=alexnet(True) 11 | input=Variable(torch.ones([1,3,226,226])) 12 | pytorch_to_caffe.trans_net(net,input,name) 13 | pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name)) 14 | pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name)) -------------------------------------------------------------------------------- /example/inceptionv3_pytorch_to_caffe.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.path.insert(0,'.') 3 | import torch 4 | from torch.autograd import Variable 5 | from torchvision.models.inception import inception_v3 6 | import pytorch_to_caffe 7 | 8 | if __name__=='__main__': 9 | name='inception_v3' 10 | net=inception_v3(True,transform_input=False) 11 | net.eval() 12 | input=torch.ones([1,3,299,299]) 13 | pytorch_to_caffe.trans_net(net,input,name) 14 | pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name)) 15 | pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name)) -------------------------------------------------------------------------------- /example/mobilenetV2_pytorch_to_caffe.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | sys.path.insert(0, '.') 4 | import torch 5 | from torch.autograd import Variable 6 | from torchvision.models import resnet 7 | import pytorch_to_caffe 8 | from MobileNetV2 import MobileNetV2 9 | 10 | if __name__ == '__main__': 11 | name = 'MobileNetV2' 12 | net= MobileNetV2() 13 | checkpoint = torch.load("/home/shining/Downloads/mobilenet_v2.pth.tar") 14 | 15 | net.load_state_dict(checkpoint) 16 | net.eval() 17 | input = torch.ones([1, 3, 224, 224]) 18 | # input=torch.ones([1,3,224,224]) 19 | pytorch_to_caffe.trans_net(net, input, name) 20 | pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name)) 21 | pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name)) -------------------------------------------------------------------------------- /example/resnet_pytorch_2_caffe.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.path.insert(0,'.') 3 | import torch 4 | from torch.autograd import Variable 5 | from torchvision.models import resnet 6 | import pytorch_to_caffe 7 | 8 | if __name__=='__main__': 9 | name='resnet18' 10 | resnet18=resnet.resnet18() 11 | checkpoint = torch.load("/home/shining/Downloads/resnet18-5c106cde.pth") 12 | 13 | resnet18.load_state_dict(checkpoint) 14 | resnet18.eval() 15 | input=torch.ones([1,3,224,224]) 16 | #input=torch.ones([1,3,224,224]) 17 | pytorch_to_caffe.trans_net(resnet18,input,name) 18 | pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name)) 19 | pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name)) -------------------------------------------------------------------------------- /example/resnet_pytorch_analysis_example.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torchvision.models import resnet 4 | import pytorch_analyser 5 | 6 | if __name__=='__main__': 7 | resnet18=resnet.resnet18() 8 | input_tensor=torch.ones(1,3,224,224) 9 | blob_dict, tracked_layers=pytorch_analyser.analyse(resnet18,input_tensor) 10 | pytorch_analyser.save_csv(tracked_layers,'/tmp/analysis.csv') 11 | 12 | -------------------------------------------------------------------------------- /example/verify_deploy.py: -------------------------------------------------------------------------------- 1 | # 2018.09.06 by Shining 2 | import sys 3 | sys.path.insert(0,'/home/shining/Projects/github-projects/caffe-project/caffe/python') 4 | import caffe 5 | import torchvision.transforms as transforms 6 | import numpy as np 7 | import argparse 8 | import torch 9 | import torch.nn as nn 10 | from torch.autograd import Variable 11 | from torch.nn.parameter import Parameter 12 | from torchvision.models import resnet 13 | import time 14 | 15 | 16 | 17 | import cv2 18 | 19 | #caffe load formate 20 | def load_image_caffe(imgfile): 21 | image = caffe.io.load_image(imgfile) 22 | transformer = caffe.io.Transformer({'data': (1, 3, args.height, args.width)}) 23 | transformer.set_transpose('data', (2, 0, 1)) 24 | transformer.set_mean('data', np.array([args.meanB, args.meanG, args.meanR])) 25 | transformer.set_raw_scale('data', args.scale) 26 | transformer.set_channel_swap('data', (2, 1, 0)) 27 | 28 | image = transformer.preprocess('data', image) 29 | image = image.reshape(1, 3, args.height, args.width) 30 | return image 31 | 32 | def load_image_pytorch(imgfile): 33 | 34 | # normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], 35 | # std=[0.229, 0.224, 0.225]) 36 | # # transforms.ToTensor() 37 | # transform1 = transforms.Compose([ 38 | # transforms.ToTensor(), # range [0, 255] -> [0.0,1.0] 39 | # normalize 40 | # ] 41 | # ) 42 | # ##numpy.ndarray 43 | # img = cv2.imread(imgfile)# 读取图像 44 | # img = cv2.resize(img,(224,244)) 45 | # img1 = transform1(img) # 归一化到 [0.0,1.0] 46 | # print("img1 = ",img1) 47 | 48 | img = np.ones([1,3,args.height, args.width]) 49 | # 转化为numpy.ndarray并显示 50 | return img 51 | 52 | 53 | 54 | def forward_pytorch(weightfile, image): 55 | net=resnet.resnet18() 56 | checkpoint = torch.load(weightfile) 57 | net.load_state_dict(checkpoint) 58 | if args.cuda: 59 | net.cuda() 60 | print(net) 61 | net.eval() 62 | image = torch.from_numpy(image) 63 | if args.cuda: 64 | image = Variable(image.cuda()) 65 | else: 66 | image = Variable(image) 67 | t0 = time.time() 68 | blobs = net.forward(image) 69 | #print(blobs.data.numpy().flatten()) 70 | t1 = time.time() 71 | return t1-t0, blobs, net.parameters() 72 | 73 | # Reference from: 74 | def forward_caffe(protofile, weightfile, image): 75 | if args.cuda: 76 | caffe.set_device(0) 77 | caffe.set_mode_gpu() 78 | else: 79 | caffe.set_mode_cpu() 80 | net = caffe.Net(protofile, weightfile, caffe.TEST) 81 | net.blobs['blob1'].reshape(1, 3, args.height, args.width) 82 | net.blobs['blob1'].data[...] = image 83 | t0 = time.time() 84 | output = net.forward() 85 | t1 = time.time() 86 | return t1-t0, net.blobs, net.params 87 | 88 | if __name__ == '__main__': 89 | parser = argparse.ArgumentParser(description='convert caffe to pytorch') 90 | parser.add_argument('--protofile', default='/home/shining/Projects/github-projects/pytorch-project/nn_tools/MGN.prototxt', type=str) 91 | parser.add_argument('--weightfile', default='/home/shining/Projects/github-projects/pytorch-project/nn_tools/MGN.caffemodel', type=str) 92 | parser.add_argument('--model', default="/home/shining/Projects/github-projects/pytorch-project/nn_tools/model_100.pt", type=str) 93 | parser.add_argument('--imgfile', default='/home/shining/Projects/github-projects/pytorch-project/nn_tools/001763.jpg', type=str) 94 | parser.add_argument('--height', default=384, type=int) 95 | parser.add_argument('--width', default=128, type=int) 96 | parser.add_argument('--meanB', default=104, type=float) 97 | parser.add_argument('--meanG', default=117, type=float) 98 | parser.add_argument('--meanR', default=123, type=float) 99 | parser.add_argument('--scale', default=255, type=float) 100 | parser.add_argument('--synset_words', default='', type=str) 101 | parser.add_argument('--cuda', action='store_true', help='enables cuda') 102 | 103 | args = parser.parse_args() 104 | print(args) 105 | 106 | 107 | protofile = args.protofile 108 | weightfile = args.weightfile 109 | imgfile = args.imgfile 110 | 111 | image = load_image_pytorch(imgfile) 112 | time_pytorch, pytorch_blobs, pytorch_models,out_Tensor_pytorch = forward_pytorch(args.model, image) 113 | time_caffe, caffe_blobs, caffe_params,out_Tensor_caffe = forward_caffe(protofile, weightfile, image) 114 | 115 | print('pytorch forward time %d', time_pytorch) 116 | print('caffe forward time %d', time_caffe) 117 | 118 | print('------------ Output Difference ------------') 119 | blob_name = "cat_blob1" 120 | if args.cuda: 121 | pytorch_data = pytorch_blobs.data.cpu().numpy().flatten() 122 | else: 123 | pytorch_data = pytorch_blobs.data.numpy().flatten() 124 | caffe_data = caffe_blobs[blob_name].data[0][...].flatten() 125 | diff = abs(pytorch_data - caffe_data).sum() 126 | print('%-30s pytorch_shape: %-20s caffe_shape: %-20s output_diff: %f' % (blob_name, pytorch_data.shape, caffe_data.shape, diff/pytorch_data.size)) 127 | -------------------------------------------------------------------------------- /example/vgg19_pytorch_to_caffe.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.path.insert(0,'.') 3 | import torch 4 | from torch.autograd import Variable 5 | from torchvision.models.vgg import vgg11_bn 6 | import pytorch_to_caffe 7 | 8 | if __name__=='__main__': 9 | name='vgg11_bn' 10 | net=vgg11_bn(True) 11 | input=Variable(torch.ones([1,3,224,224])) 12 | pytorch_to_caffe.trans_net(net,input,name) 13 | pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name)) 14 | pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name)) -------------------------------------------------------------------------------- /funcs.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import numpy as np 3 | import time 4 | 5 | class Logger(): 6 | def __init__(self,file_name=None,show=True): 7 | self.show=show 8 | self.file_name=file_name 9 | 10 | def __call__(self,str): 11 | str='%s '%(time.strftime('%H:%M:%S'),)+str 12 | if self.file_name: 13 | with open(self.file_name,'a+') as f: 14 | f.write(str+'\n') 15 | if self.show: 16 | print(str) 17 | 18 | def intersect(box_a, box_b): 19 | max_xy = np.minimum(box_a[:, 2:], box_b[2:]) 20 | min_xy = np.maximum(box_a[:, :2], box_b[:2]) 21 | inter = np.clip((max_xy - min_xy), a_min=0, a_max=np.inf) 22 | return inter[:, 0] * inter[:, 1] 23 | 24 | def get_iou(box_a, box_b): 25 | """Compute the jaccard overlap of two sets of boxes. The jaccard overlap 26 | is simply the intersection over union of two boxes. 27 | E.g.: 28 | A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B) 29 | The box should be [x1,y1,x2,y2] 30 | Args: 31 | box_a: Single numpy bounding box, Shape: [4] or Multiple bounding boxes, Shape: [num_boxes,4] 32 | box_b: Single numpy bounding box, Shape: [4] 33 | Return: 34 | jaccard overlap: Shape: [box_a.shape[0], box_a.shape[1]] 35 | """ 36 | if box_a.ndim==1: 37 | box_a=box_a.reshape([1,-1]) 38 | inter = intersect(box_a, box_b) 39 | area_a = ((box_a[:, 2]-box_a[:, 0]) * 40 | (box_a[:, 3]-box_a[:, 1])) # [A,B] 41 | area_b = ((box_b[2]-box_b[0]) * 42 | (box_b[3]-box_b[1])) # [A,B] 43 | union = area_a + area_b - inter 44 | return inter / union # [A,B] 45 | 46 | def nms(bboxs,scores,thresh): 47 | """ 48 | The box should be [x1,y1,x2,y2] 49 | :param bboxs: multiple bounding boxes, Shape: [num_boxes,4] 50 | :param scores: The score for the corresponding box 51 | :return: keep inds 52 | """ 53 | if len(bboxs)==0: 54 | return [] 55 | order=scores.argsort()[::-1] 56 | keep=[] 57 | while order.size>0: 58 | i=order[0] 59 | keep.append(i) 60 | ious=get_iou(bboxs[order],bboxs[i]) 61 | order=order[ious<=thresh] 62 | return keep 63 | -------------------------------------------------------------------------------- /model/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | from importlib import import_module 3 | 4 | import torch 5 | import torch.nn as nn 6 | 7 | class Model(nn.Module): 8 | def __init__(self, args, ckpt): 9 | super(Model, self).__init__() 10 | print('[INFO] Making model...') 11 | 12 | self.device = torch.device('cpu' if args.cpu else 'cuda') 13 | self.nGPU = args.nGPU 14 | self.save_models = args.save_models 15 | 16 | module = import_module('model.' + args.model.lower()) 17 | self.model = module.make_model(args).to(self.device) 18 | 19 | if not args.cpu and args.nGPU > 1: 20 | self.model = nn.DataParallel(self.model, range(args.nGPU)) 21 | 22 | self.load( 23 | ckpt.dir, 24 | pre_train=args.pre_train, 25 | resume=args.resume, 26 | cpu=args.cpu 27 | ) 28 | print(self.model, file=ckpt.log_file) 29 | 30 | 31 | def forward(self, x): 32 | return self.model(x) 33 | 34 | def get_model(self): 35 | if self.nGPU == 1: 36 | return self.model 37 | else: 38 | return self.model.module 39 | 40 | def save(self, apath, epoch, is_best=False): 41 | target = self.get_model() 42 | torch.save( 43 | target.state_dict(), 44 | os.path.join(apath, 'model', 'model_latest.pt') 45 | ) 46 | if is_best: 47 | torch.save( 48 | target.state_dict(), 49 | os.path.join(apath, 'model', 'model_best.pt') 50 | ) 51 | 52 | if self.save_models: 53 | torch.save( 54 | target.state_dict(), 55 | os.path.join(apath, 'model', 'model_{}.pt'.format(epoch)) 56 | ) 57 | 58 | def load(self, apath, pre_train='', resume=-1, cpu=False): 59 | if cpu: 60 | kwargs = {'map_location': lambda storage, loc: storage} 61 | else: 62 | kwargs = {} 63 | 64 | if resume == -1: 65 | self.get_model().load_state_dict( 66 | torch.load( 67 | os.path.join(apath, 'model', 'model_latest.pt'), 68 | **kwargs 69 | ), 70 | strict=False 71 | ) 72 | elif resume == 0: 73 | if pre_train != '': 74 | print('Loading model from {}'.format(pre_train)) 75 | self.get_model().load_state_dict( 76 | torch.load(pre_train, **kwargs), 77 | strict=False 78 | ) 79 | else: 80 | self.get_model().load_state_dict( 81 | torch.load( 82 | '/home/wdd/Work/Pytorch/MGN-pytorch-master/experiment/test815/model/model_100.pt', 83 | **kwargs 84 | ), 85 | strict=False 86 | ) -------------------------------------------------------------------------------- /model/__pycache__/__init__.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/model/__pycache__/__init__.cpython-35.pyc -------------------------------------------------------------------------------- /model/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/model/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /model/__pycache__/mgn.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/model/__pycache__/mgn.cpython-35.pyc -------------------------------------------------------------------------------- /model/__pycache__/mgn.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxradon/PytorchToCaffe/7aa8422eb2b5ac90ceb47e7658814241aa7037b2/model/__pycache__/mgn.cpython-36.pyc -------------------------------------------------------------------------------- /model/mgn.py: -------------------------------------------------------------------------------- 1 | import copy 2 | 3 | import torch 4 | import torch.nn.functional as F 5 | 6 | from torchvision.models.resnet import Bottleneck,resnet50 7 | import torch.nn as nn 8 | 9 | 10 | 11 | 12 | 13 | 14 | def make_model(args): 15 | return MGN(args) 16 | 17 | class MGN(nn.Module): 18 | def __init__(self, args): 19 | super(MGN, self).__init__() 20 | num_classes = args.num_classes 21 | 22 | resnet = resnet50(pretrained=True) 23 | 24 | self.backone = nn.Sequential( 25 | resnet.conv1, 26 | resnet.bn1, 27 | resnet.relu, 28 | resnet.maxpool, 29 | resnet.layer1, 30 | resnet.layer2, 31 | resnet.layer3[0], 32 | ) 33 | 34 | res_conv4 = nn.Sequential(*resnet.layer3[1:]) 35 | # res_conv4 = nn.Sequential(resnet.layer3[1],resnet.layer3[2],resnet.layer3[3],resnet.layer3[4]，resnet.layer3[5]) 36 | 37 | res_g_conv5 = resnet.layer4 38 | 39 | res_p_conv5 = nn.Sequential( 40 | Bottleneck(1024, 512, downsample=nn.Sequential(nn.Conv2d(1024, 2048, 1, bias=False), nn.BatchNorm2d(2048))), 41 | Bottleneck(2048, 512), 42 | Bottleneck(2048, 512)) 43 | res_p_conv5.load_state_dict(resnet.layer4.state_dict()) 44 | 45 | self.p1 = nn.Sequential(copy.deepcopy(res_conv4), copy.deepcopy(res_g_conv5)) 46 | self.p2 = nn.Sequential(copy.deepcopy(res_conv4), copy.deepcopy(res_p_conv5)) 47 | self.p3 = nn.Sequential(copy.deepcopy(res_conv4), copy.deepcopy(res_p_conv5)) 48 | 49 | if args.pool == 'max': 50 | pool2d = nn.MaxPool2d 51 | elif args.pool == 'avg': 52 | pool2d = nn.AvgPool2d 53 | else: 54 | raise Exception() 55 | 56 | 57 | self.maxpool_zg_p1 = pool2d(kernel_size=(12, 4)) 58 | self.maxpool_zg_p2 = pool2d(kernel_size=(24, 8)) 59 | self.maxpool_zg_p3 = pool2d(kernel_size=(24, 8)) 60 | self.maxpool_zp2 = pool2d(kernel_size=(12, 8)) 61 | self.maxpool_zp3 = pool2d(kernel_size=(8, 8)) 62 | 63 | # self.maxpool_zg_p1 = nn.AvgPool2d(kernel_size=(12 * 2, 4 * 2)) 64 | # self.maxpool_zg_p2 = nn.AvgPool2d(kernel_size=(24, 8)) 65 | # self.maxpool_zg_p3 = nn.AvgPool2d(kernel_size=(24, 8)) 66 | # self.maxpool_zp2 = nn.MaxPool2d(kernel_size=(12, 8)) 67 | # self.maxpool_zp3 = nn.MaxPool2d(kernel_size=(8, 8)) 68 | 69 | reduction = nn.Sequential(nn.Conv2d(2048, args.feats, 1, bias=False), nn.BatchNorm2d(args.feats), nn.ReLU()) 70 | 71 | self._init_reduction(reduction) 72 | self.reduction_0 = copy.deepcopy(reduction) 73 | self.reduction_1 = copy.deepcopy(reduction) 74 | self.reduction_2 = copy.deepcopy(reduction) 75 | self.reduction_3 = copy.deepcopy(reduction) 76 | self.reduction_4 = copy.deepcopy(reduction) 77 | self.reduction_5 = copy.deepcopy(reduction) 78 | self.reduction_6 = copy.deepcopy(reduction) 79 | self.reduction_7 = copy.deepcopy(reduction) 80 | 81 | #self.fc_id_2048_0 = nn.Linear(2048, num_classes) 82 | #=============code autor ============================= 83 | self.fc_id_2048_0 = nn.Linear(args.feats, num_classes) 84 | self.fc_id_2048_1 = nn.Linear(args.feats, num_classes) 85 | self.fc_id_2048_2 = nn.Linear(args.feats, num_classes) 86 | #===================================================== 87 | 88 | 89 | #============papers autor======================= 90 | # self.fc_id_2048_0 = nn.Linear(2048, num_classes) 91 | # self.fc_id_2048_1 = nn.Linear(2048, num_classes) 92 | # self.fc_id_2048_2 = nn.Linear(2048, num_classes) 93 | #================================================ 94 | 95 | self.fc_id_256_1_0 = nn.Linear(args.feats, num_classes) 96 | self.fc_id_256_1_1 = nn.Linear(args.feats, num_classes) 97 | self.fc_id_256_2_0 = nn.Linear(args.feats, num_classes) 98 | self.fc_id_256_2_1 = nn.Linear(args.feats, num_classes) 99 | self.fc_id_256_2_2 = nn.Linear(args.feats, num_classes) 100 | 101 | self._init_fc(self.fc_id_2048_0) 102 | self._init_fc(self.fc_id_2048_1) 103 | self._init_fc(self.fc_id_2048_2) 104 | 105 | self._init_fc(self.fc_id_256_1_0) 106 | self._init_fc(self.fc_id_256_1_1) 107 | self._init_fc(self.fc_id_256_2_0) 108 | self._init_fc(self.fc_id_256_2_1) 109 | self._init_fc(self.fc_id_256_2_2) 110 | 111 | @staticmethod 112 | def _init_reduction(reduction): 113 | # conv 114 | nn.init.kaiming_normal_(reduction[0].weight, mode='fan_in') 115 | #nn.init.constant_(reduction[0].bias, 0.) 116 | 117 | # bn 118 | nn.init.normal_(reduction[1].weight, mean=1., std=0.02) 119 | nn.init.constant_(reduction[1].bias, 0.) 120 | 121 | @staticmethod 122 | def _init_fc(fc): 123 | nn.init.kaiming_normal_(fc.weight, mode='fan_out') 124 | #nn.init.normal_(fc.weight, std=0.001) 125 | nn.init.constant_(fc.bias, 0.) 126 | 127 | def forward(self, x): 128 | 129 | x = self.backone(x) 130 | 131 | p1 = self.p1(x) 132 | p2 = self.p2(x) 133 | p3 = self.p3(x) 134 | 135 | zg_p1 = self.maxpool_zg_p1(p1) 136 | zg_p2 = self.maxpool_zg_p2(p2) 137 | zg_p3 = self.maxpool_zg_p3(p3) 138 | 139 | zp2 = self.maxpool_zp2(p2) 140 | #z0_p2 = zp2[:, :, 0:1, :] 141 | z0_p2,z1_p2 = torch.split(zp2, 1 ,2) 142 | # indices = torch.tensor([0]) 143 | # z0_p2 = torch.index_select(zp2, 2 , indices) 144 | #z1_p2 = zp2[:, :, 1:2, :] 145 | # indices = torch.tensor([1]) 146 | # z1_p2 = torch.index_select(zp2, 2 , indices) 147 | 148 | 149 | zp3 = self.maxpool_zp3(p3) 150 | 151 | #z0_p3 = zp3[:, :, 0:1, :] 152 | z0_p3,z1_p3,z2_p3 = torch.split(zp3,1,2) 153 | 154 | # indices = torch.tensor([0]) 155 | # z0_p3 = torch.index_select(zp3, 2 , indices) 156 | #z1_p3 = zp3[:, :, 1:2, :] 157 | # indices = torch.tensor([1]) 158 | # z1_p3 = torch.index_select(zp3, 2 , indices) 159 | # #z2_p3 = zp3[:, :, 2:3, :] 160 | # indices = torch.tensor([2]) 161 | # z2_p3 = torch.index_select(zp3, 2 , indices) 162 | 163 | fg_p1 = self.reduction_0(zg_p1) 164 | fg_p2 = self.reduction_1(zg_p2) 165 | fg_p3 = self.reduction_2(zg_p3) 166 | f0_p2 = self.reduction_3(z0_p2) 167 | f1_p2 = self.reduction_4(z1_p2) 168 | f0_p3 = self.reduction_5(z0_p3) 169 | f1_p3 = self.reduction_6(z1_p3) 170 | f2_p3 = self.reduction_7(z2_p3) 171 | predict = torch.cat([fg_p1, fg_p2, fg_p3, f0_p2, f1_p2, f0_p3, f1_p3, f2_p3], dim=1) 172 | return predict 173 | 174 | # fg_p1 = self.reduction_0(zg_p1).squeeze(dim=3).squeeze(dim=2) 175 | # fg_p2 = self.reduction_1(zg_p2).squeeze(dim=3).squeeze(dim=2) 176 | # fg_p3 = self.reduction_2(zg_p3).squeeze(dim=3).squeeze(dim=2) 177 | # f0_p2 = self.reduction_3(z0_p2).squeeze(dim=3).squeeze(dim=2) 178 | # f1_p2 = self.reduction_4(z1_p2).squeeze(dim=3).squeeze(dim=2) 179 | # f0_p3 = self.reduction_5(z0_p3).squeeze(dim=3).squeeze(dim=2) 180 | # f1_p3 = self.reduction_6(z1_p3).squeeze(dim=3).squeeze(dim=2) 181 | # f2_p3 = self.reduction_7(z2_p3).squeeze(dim=3).squeeze(dim=2) 182 | 183 | # ''' 184 | # #================papers autor=============================== 185 | # l_p1 = self.fc_id_2048_0(zg_p1.squeeze(dim=3).squeeze(dim=2)) 186 | # l_p2 = self.fc_id_2048_1(zg_p2.squeeze(dim=3).squeeze(dim=2)) 187 | # l_p3 = self.fc_id_2048_2(zg_p3.squeeze(dim=3).squeeze(dim=2)) 188 | # ''' 189 | 190 | # l_p1 = self.fc_id_2048_0(fg_p1) 191 | # l_p2 = self.fc_id_2048_1(fg_p2) 192 | # l_p3 = self.fc_id_2048_2(fg_p3) 193 | 194 | # l0_p2 = self.fc_id_256_1_0(f0_p2) 195 | # l1_p2 = self.fc_id_256_1_1(f1_p2) 196 | # l0_p3 = self.fc_id_256_2_0(f0_p3) 197 | # l1_p3 = self.fc_id_256_2_1(f1_p3) 198 | # l2_p3 = self.fc_id_256_2_2(f2_p3) 199 | 200 | 201 | # predict = fg_p1 202 | # predict = torch.cat([fg_p1, fg_p2, fg_p3],dim=1) 203 | # predict = (fg_p1+fg_p2+fg_p3+f0_p2+f1_p2+f0_p3+f1_p3+f2_p3)/8 204 | # predict = (fg_p1+fg_p2+fg_p3)/3 205 | #return predict, fg_p1, fg_p2, fg_p3, l_p1, l_p2, l_p3, l0_p2, l1_p2, l0_p3, l1_p3, l2_p3 206 | 207 | 208 | 209 | 210 | -------------------------------------------------------------------------------- /model/resnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import math 3 | import torch.utils.model_zoo as model_zoo 4 | 5 | 6 | __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 7 | 'resnet152'] 8 | 9 | 10 | model_urls = { 11 | 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 12 | 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 13 | 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 14 | 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 15 | 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 16 | } 17 | 18 | 19 | def conv3x3(in_planes, out_planes, stride=1): 20 | """3x3 convolution with padding""" 21 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 22 | padding=1, bias=False) 23 | 24 | 25 | class BasicBlock(nn.Module): 26 | expansion = 1 27 | 28 | def __init__(self, inplanes, planes, stride=1, downsample=None): 29 | super(BasicBlock, self).__init__() 30 | self.conv1 = conv3x3(inplanes, planes, stride) 31 | self.bn1 = nn.BatchNorm2d(planes) 32 | self.relu = nn.ReLU(inplace=True) 33 | self.conv2 = conv3x3(planes, planes) 34 | self.bn2 = nn.BatchNorm2d(planes) 35 | self.downsample = downsample 36 | self.stride = stride 37 | 38 | def forward(self, x): 39 | residual = x 40 | 41 | out = self.conv1(x) 42 | out = self.bn1(out) 43 | out = self.relu(out) 44 | 45 | out = self.conv2(out) 46 | out = self.bn2(out) 47 | 48 | if self.downsample is not None: 49 | residual = self.downsample(x) 50 | 51 | out += residual 52 | out = self.relu(out) 53 | 54 | return out 55 | 56 | 57 | class Bottleneck(nn.Module): 58 | expansion = 4 59 | 60 | def __init__(self, inplanes, planes, stride=1, downsample=None): 61 | super(Bottleneck, self).__init__() 62 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 63 | self.bn1 = nn.BatchNorm2d(planes) 64 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, 65 | padding=1, bias=False) 66 | self.bn2 = nn.BatchNorm2d(planes) 67 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 68 | self.bn3 = nn.BatchNorm2d(planes * 4) 69 | self.relu = nn.ReLU(inplace=True) 70 | self.downsample = downsample 71 | self.stride = stride 72 | 73 | def forward(self, x): 74 | residual = x 75 | 76 | out = self.conv1(x) 77 | out = self.bn1(out) 78 | out = self.relu(out) 79 | 80 | out = self.conv2(out) 81 | out = self.bn2(out) 82 | out = self.relu(out) 83 | 84 | out = self.conv3(out) 85 | out = self.bn3(out) 86 | 87 | if self.downsample is not None: 88 | residual = self.downsample(x) 89 | 90 | out += residual 91 | out = self.relu(out) 92 | 93 | return out 94 | 95 | 96 | class ResNet(nn.Module): 97 | 98 | def __init__(self, block, layers, num_classes=1000): 99 | self.inplanes = 64 100 | super(ResNet, self).__init__() 101 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 102 | bias=False) 103 | self.bn1 = nn.BatchNorm2d(64) 104 | self.relu = nn.ReLU(inplace=True) 105 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 106 | self.layer1 = self._make_layer(block, 64, layers[0]) 107 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 108 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 109 | self.layer4 = self._make_layer(block, 512, layers[3], stride=1) 110 | self.avgpool = nn.AvgPool2d(7, stride=1) 111 | self.fc = nn.Linear(512 * block.expansion, num_classes) 112 | 113 | for m in self.modules(): 114 | if isinstance(m, nn.Conv2d): 115 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 116 | m.weight.data.normal_(0, math.sqrt(2. / n)) 117 | elif isinstance(m, nn.BatchNorm2d): 118 | m.weight.data.fill_(1) 119 | m.bias.data.zero_() 120 | 121 | def _make_layer(self, block, planes, blocks, stride=1): 122 | downsample = None 123 | if stride != 1 or self.inplanes != planes * block.expansion: 124 | downsample = nn.Sequential( 125 | nn.Conv2d(self.inplanes, planes * block.expansion, 126 | kernel_size=1, stride=stride, bias=False), 127 | nn.BatchNorm2d(planes * block.expansion), 128 | ) 129 | 130 | layers = [] 131 | layers.append(block(self.inplanes, planes, stride, downsample)) 132 | self.inplanes = planes * block.expansion 133 | for i in range(1, blocks): 134 | layers.append(block(self.inplanes, planes)) 135 | 136 | return nn.Sequential(*layers) 137 | 138 | def forward(self, x): 139 | x = self.conv1(x) 140 | x = self.bn1(x) 141 | x = self.relu(x) 142 | x = self.maxpool(x) 143 | 144 | x = self.layer1(x) 145 | x = self.layer2(x) 146 | x = self.layer3(x) 147 | x = self.layer4(x) 148 | 149 | x = self.avgpool(x) 150 | x = x.view(x.size(0), -1) 151 | x = self.fc(x) 152 | 153 | return x 154 | 155 | 156 | def resnet18(pretrained=False, **kwargs): 157 | """Constructs a ResNet-18 model. 158 | 159 | Args: 160 | pretrained (bool): If True, returns a model pre-trained on ImageNet 161 | """ 162 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 163 | if pretrained: 164 | model.load_state_dict(model_zoo.load_url(model_urls['resnet18'])) 165 | return model 166 | 167 | 168 | def resnet34(pretrained=False, **kwargs): 169 | """Constructs a ResNet-34 model. 170 | 171 | Args: 172 | pretrained (bool): If True, returns a model pre-trained on ImageNet 173 | """ 174 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 175 | if pretrained: 176 | model.load_state_dict(model_zoo.load_url(model_urls['resnet34'])) 177 | return model 178 | 179 | 180 | def resnet50(pretrained=False, **kwargs): 181 | """Constructs a ResNet-50 model. 182 | 183 | Args: 184 | pretrained (bool): If True, returns a model pre-trained on ImageNet 185 | """ 186 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 187 | if pretrained: 188 | model.load_state_dict(model_zoo.load_url(model_urls['resnet50'])) 189 | return model 190 | 191 | 192 | def resnet101(pretrained=False, **kwargs): 193 | """Constructs a ResNet-101 model. 194 | 195 | Args: 196 | pretrained (bool): If True, returns a model pre-trained on ImageNet 197 | """ 198 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 199 | if pretrained: 200 | model.load_state_dict(model_zoo.load_url(model_urls['resnet101'])) 201 | return model 202 | 203 | 204 | def resnet152(pretrained=False, **kwargs): 205 | """Constructs a ResNet-152 model. 206 | 207 | Args: 208 | pretrained (bool): If True, returns a model pre-trained on ImageNet 209 | """ 210 | model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) 211 | if pretrained: 212 | model.load_state_dict(model_zoo.load_url(model_urls['resnet152'])) 213 | return model 214 | -------------------------------------------------------------------------------- /pytorch_analyser.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | from __future__ import absolute_import 3 | import torch 4 | import argparse 5 | import sys,os 6 | from analysis.PytorchA import analyse 7 | from analysis.utils import save_csv 8 | from torch.autograd import Variable 9 | import torch.nn as nn 10 | 11 | """ 12 | Supporting analyse the inheritors of torch.nn.Moudule class. 13 | 14 | Command：`pytorch_analyser.py [-h] [--out OUT] [--class_args ARGS] path class_name shape` 15 | - The path is the python file path which contaning your class. 16 | - The class_name is the class name in your python file. 17 | - The shape is the input shape of the network(split by comma `,`), in pytorch image shape should be: batch_size, channel, image_height, image_width. 18 | - The out (optinal) is path to save the csv file, default is '/tmp/pytorch_analyse.csv'. 19 | - The class_args (optional) is the args to init the class in python file, default is empty. 20 | 21 | For example `python pytorch_analyser.py tmp/pytorch_analysis_test.py ResNet218 1,3,224,224` 22 | """ 23 | 24 | 25 | if __name__=="__main__": 26 | parser=argparse.ArgumentParser() 27 | parser.add_argument('--path',help='python file location',type=str) 28 | parser.add_argument('--name',help='the class name or instance name in your python file',type=str) 29 | parser.add_argument('--shape',help='input shape of the network(split by comma `,`), image shape should be: batch,c,h,w',type=str) 30 | parser.add_argument('--out',help='path to save the csv file',default='/tmp/pytorch_analyse.csv',type=str) 31 | parser.add_argument('--class_args',help='args to init the class in python file',default='',type=str) 32 | 33 | args=parser.parse_args() 34 | path,filename=os.path.split(args.path) 35 | filename=os.path.splitext(filename)[0] 36 | sys.path.insert(0,path) 37 | exec('from %s import %s as Net'%(filename,args.name)) 38 | if isinstance(Net, nn.Module): 39 | net=Net 40 | elif issubclass(Net,nn.Module): 41 | net=Net(*args.class_args.split()) 42 | else: 43 | assert("Error, The Net is not a instance of nn.Module or subclass of nn.Module") 44 | shape = [int(i) for i in args.shape.split(',')] 45 | x = Variable(torch.rand(shape)) 46 | blob_dict, layers = analyse(net, x) 47 | save_csv(layers, args.out) 48 | 49 | 50 | -------------------------------------------------------------------------------- /pytorch_to_caffe.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import traceback 4 | from Caffe import caffe_net 5 | import torch.nn.functional as F 6 | from torch.autograd import Variable 7 | from Caffe import layer_param 8 | from torch.nn.modules.utils import _pair 9 | import numpy as np 10 | 11 | """ 12 | How to support a new layer type: 13 | layer_name=log.add_layer(layer_type_name) 14 | top_blobs=log.add_blobs() 15 | layer=caffe_net.Layer_param(xxx) 16 | 17 | [] 18 | log.cnet.add_layer(layer) 19 | 20 | Please MUTE the inplace operations to avoid not find in graph 21 | 22 | 注意：只有torch.nn.functional中的函数才能转换为caffe中的层 23 | """ 24 | 25 | # TODO: support the inplace output of the layers 26 | 27 | class Blob_LOG(): 28 | def __init__(self): 29 | self.data={} 30 | def __setitem__(self, key, value): 31 | self.data[key]=value 32 | def __getitem__(self, key): 33 | return self.data[key] 34 | def __len__(self): 35 | return len(self.data) 36 | 37 | NET_INITTED=False 38 | 39 | # 转换原理解析：通过记录 40 | class TransLog(object): 41 | def __init__(self): 42 | """ 43 | doing init() with inputs Variable before using it 44 | """ 45 | self.layers={} 46 | self.detail_layers={} 47 | self.detail_blobs={} 48 | self._blobs=Blob_LOG() 49 | self._blobs_data=[] 50 | self.cnet=caffe_net.Caffemodel('') 51 | self.debug=True 52 | 53 | def init(self,inputs): 54 | """ 55 | :param inputs: is a list of input variables 56 | """ 57 | self.add_blobs(inputs) 58 | def add_layer(self,name='layer'): 59 | if name in self.layers: 60 | return self.layers[name] 61 | if name not in self.detail_layers.keys(): 62 | self.detail_layers[name] =0 63 | self.detail_layers[name] +=1 64 | name='{}{}'.format(name,self.detail_layers[name]) 65 | self.layers[name]=name 66 | if self.debug: 67 | print("{} was added to layers".format(self.layers[name])) 68 | return self.layers[name] 69 | 70 | def add_blobs(self, blobs,name='blob',with_num=True): 71 | rst=[] 72 | for blob in blobs: 73 | self._blobs_data.append(blob) # to block the memory address be rewrited 74 | blob_id=int(id(blob)) 75 | if name not in self.detail_blobs.keys(): 76 | self.detail_blobs[name] =0 77 | self.detail_blobs[name] +=1 78 | if with_num: 79 | rst.append('{}{}'.format(name,self.detail_blobs[name])) 80 | else: 81 | rst.append('{}'.format(name)) 82 | if self.debug: 83 | print("{}:{} was added to blobs".format(blob_id,rst[-1])) 84 | # print('Add blob {} : {}'.format(rst[-1].center(21),blob.size())) 85 | self._blobs[blob_id]=rst[-1] 86 | return rst 87 | def blobs(self, var): 88 | var=id(var) 89 | # if self.debug: 90 | # print("{}:{} getting".format(var, self._blobs[var])) 91 | try: 92 | return self._blobs[var] 93 | except: 94 | print("WARNING: CANNOT FOUND blob {}".format(var)) 95 | return None 96 | 97 | log=TransLog() 98 | 99 | layer_names={} 100 | def _conv2d(raw,input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): 101 | print('conv: ',log.blobs(input)) 102 | x=raw(input,weight,bias,stride,padding,dilation,groups) 103 | name=log.add_layer(name='conv') 104 | log.add_blobs([x],name='conv_blob') 105 | layer=caffe_net.Layer_param(name=name, type='Convolution', 106 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 107 | layer.conv_param(x.size()[1],weight.size()[2:],stride=_pair(stride), 108 | pad=_pair(padding),dilation=_pair(dilation),bias_term=bias is not None,groups=groups) 109 | if bias is not None: 110 | layer.add_data(weight.cpu().data.numpy(),bias.cpu().data.numpy()) 111 | else: 112 | layer.param.convolution_param.bias_term=False 113 | layer.add_data(weight.cpu().data.numpy()) 114 | log.cnet.add_layer(layer) 115 | return x 116 | 117 | def _conv_transpose2d(raw,input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1): 118 | x=raw(input, weight, bias, stride, padding, output_padding, groups, dilation) 119 | name=log.add_layer(name='conv_transpose') 120 | log.add_blobs([x],name='conv_transpose_blob') 121 | layer=caffe_net.Layer_param(name=name, type='Deconvolution', 122 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 123 | layer.conv_param(x.size()[1],weight.size()[2:],stride=_pair(stride), 124 | pad=_pair(padding),dilation=_pair(dilation),bias_term=bias is not None, groups = groups) 125 | if bias is not None: 126 | layer.add_data(weight.cpu().data.numpy(),bias.cpu().data.numpy()) 127 | else: 128 | layer.param.convolution_param.bias_term=False 129 | layer.add_data(weight.cpu().data.numpy()) 130 | log.cnet.add_layer(layer) 131 | return x 132 | 133 | def _linear(raw,input, weight, bias=None): 134 | x=raw(input,weight,bias) 135 | layer_name=log.add_layer(name='fc') 136 | top_blobs=log.add_blobs([x],name='fc_blob') 137 | layer=caffe_net.Layer_param(name=layer_name,type='InnerProduct', 138 | bottom=[log.blobs(input)],top=top_blobs) 139 | layer.fc_param(x.size()[1],has_bias=bias is not None) 140 | if bias is not None: 141 | layer.add_data(weight.cpu().data.numpy(),bias.cpu().data.numpy()) 142 | else: 143 | layer.add_data(weight.cpu().data.numpy()) 144 | log.cnet.add_layer(layer) 145 | return x 146 | 147 | def _split(raw,tensor, split_size, dim=0): 148 | # split in pytorch is slice in caffe 149 | x=raw(tensor, split_size, dim) 150 | layer_name=log.add_layer('split') 151 | top_blobs=log.add_blobs(x,name='split_blob') 152 | layer=caffe_net.Layer_param(name=layer_name, type='Slice', 153 | bottom=[log.blobs(tensor)], top=top_blobs) 154 | slice_num=int(np.floor(tensor.size()[dim]/split_size)) 155 | slice_param=caffe_net.pb.SliceParameter(axis=dim,slice_point=[split_size*i for i in range(1,slice_num)]) 156 | layer.param.slice_param.CopyFrom(slice_param) 157 | log.cnet.add_layer(layer) 158 | return x 159 | 160 | 161 | def _pool(type,raw,input,x,kernel_size,stride,padding,ceil_mode): 162 | # TODO dilation,ceil_mode,return indices 163 | layer_name = log.add_layer(name='{}_pool'.format(type)) 164 | top_blobs = log.add_blobs([x], name='{}_pool_blob'.format(type)) 165 | layer = caffe_net.Layer_param(name=layer_name, type='Pooling', 166 | bottom=[log.blobs(input)], top=top_blobs) 167 | # TODO w,h different kernel, stride and padding 168 | # processing ceil mode 169 | layer.pool_param(kernel_size=kernel_size, stride=kernel_size if stride is None else stride, 170 | pad=padding, type=type.upper() , ceil_mode = ceil_mode) 171 | log.cnet.add_layer(layer) 172 | if ceil_mode==False and stride is not None: 173 | oheight = (input.size()[2] - _pair(kernel_size)[0] + 2 * _pair(padding)[0]) % (_pair(stride)[0]) 174 | owidth = (input.size()[3] - _pair(kernel_size)[1] + 2 * _pair(padding)[1]) % (_pair(stride)[1]) 175 | if oheight!=0 or owidth!=0: 176 | caffe_out=raw(input, kernel_size, stride, padding, ceil_mode=True) 177 | print("WARNING: the output shape miss match at {}: " 178 | 179 | "input {} output---Pytorch:{}---Caffe:{}\n" 180 | "This is caused by the different implementation that ceil mode in caffe and the floor mode in pytorch.\n" 181 | "You can add the clip layer in caffe prototxt manually if shape mismatch error is caused in caffe. ".format(layer_name,input.size(),x.size(),caffe_out.size())) 182 | 183 | def _max_pool2d(raw,input, kernel_size, stride=None, padding=0, dilation=1, 184 | ceil_mode=False, return_indices=False): 185 | x = raw(input, kernel_size, stride, padding, dilation,ceil_mode, return_indices) 186 | _pool('max',raw,input, x, kernel_size, stride, padding,ceil_mode) 187 | return x 188 | 189 | def _avg_pool2d(raw,input, kernel_size, stride = None, padding = 0, ceil_mode = False, count_include_pad = True): 190 | x = raw(input, kernel_size, stride, padding, ceil_mode, count_include_pad) 191 | _pool('ave',raw,input, x, kernel_size, stride, padding,ceil_mode) 192 | return x 193 | 194 | def _adaptive_avg_pool2d(raw, input, output_size): 195 | x = raw(input, output_size) 196 | if isinstance(output_size, int): 197 | out_dim = output_size 198 | else: 199 | out_dim = output_size[0] 200 | tmp = max(input.shape[2], input.shape[3]) 201 | stride = tmp //out_dim 202 | kernel_size = tmp - (out_dim - 1) * stride 203 | _pool('ave', raw, input, x, kernel_size, stride, 0, False) 204 | return x 205 | 206 | def _max(raw,*args): 207 | x=raw(*args) 208 | if len(args)==1: 209 | # TODO max in one tensor 210 | assert NotImplementedError 211 | else: 212 | bottom_blobs=[] 213 | for arg in args: 214 | bottom_blobs.append(log.blobs(arg)) 215 | layer_name=log.add_layer(name='max') 216 | top_blobs=log.add_blobs([x],name='max_blob') 217 | layer=caffe_net.Layer_param(name=layer_name,type='Eltwise', 218 | bottom=bottom_blobs,top=top_blobs) 219 | layer.param.eltwise_param.operation =2 220 | log.cnet.add_layer(layer) 221 | return x 222 | 223 | def _cat(raw,inputs, dimension=0): 224 | x=raw(inputs, dimension) 225 | bottom_blobs=[] 226 | for input in inputs: 227 | bottom_blobs.append(log.blobs(input)) 228 | layer_name=log.add_layer(name='cat') 229 | top_blobs=log.add_blobs([x],name='cat_blob') 230 | layer=caffe_net.Layer_param(name=layer_name,type='Concat', 231 | bottom=bottom_blobs,top=top_blobs) 232 | layer.param.concat_param.axis =dimension 233 | log.cnet.add_layer(layer) 234 | return x 235 | 236 | def _dropout(raw,input,p=0.5, training=False, inplace=False): 237 | x=raw(input,p, training, inplace) 238 | bottom_blobs=[log.blobs(input)] 239 | layer_name=log.add_layer(name='dropout') 240 | top_blobs=log.add_blobs([x],name=bottom_blobs[0],with_num=False) 241 | layer=caffe_net.Layer_param(name=layer_name,type='Dropout', 242 | bottom=bottom_blobs,top=top_blobs) 243 | layer.param.dropout_param.dropout_ratio = p 244 | layer.param.include.extend([caffe_net.pb.NetStateRule(phase=0)]) # 1 for test, 0 for train 245 | log.cnet.add_layer(layer) 246 | return x 247 | 248 | def _threshold(raw,input, threshold, value, inplace=False): 249 | # for threshold or relu 250 | if threshold==0 and value==0: 251 | x = raw(input,threshold, value, inplace) 252 | bottom_blobs=[log.blobs(input)] 253 | name = log.add_layer(name='relu') 254 | log.add_blobs([x], name='relu_blob') 255 | layer = caffe_net.Layer_param(name=name, type='ReLU', 256 | bottom=bottom_blobs, top=[log.blobs(x)]) 257 | log.cnet.add_layer(layer) 258 | return x 259 | if value!=0: 260 | raise NotImplemented("value !=0 not implemented in caffe") 261 | x=raw(input,input, threshold, value, inplace) 262 | bottom_blobs=[log.blobs(input)] 263 | layer_name=log.add_layer(name='threshold') 264 | top_blobs=log.add_blobs([x],name='threshold_blob') 265 | layer=caffe_net.Layer_param(name=layer_name,type='Threshold', 266 | bottom=bottom_blobs,top=top_blobs) 267 | layer.param.threshold_param.threshold = threshold 268 | log.cnet.add_layer(layer) 269 | return x 270 | 271 | def _relu(raw, input, inplace=False): 272 | # for threshold or prelu 273 | x = raw(input, False) 274 | name = log.add_layer(name='relu') 275 | log.add_blobs([x], name='relu_blob') 276 | layer = caffe_net.Layer_param(name=name, type='ReLU', 277 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 278 | log.cnet.add_layer(layer) 279 | return x 280 | def _prelu(raw, input, weight): 281 | # for threshold or prelu 282 | x = raw(input, weight) 283 | bottom_blobs=[log.blobs(input)] 284 | name = log.add_layer(name='prelu') 285 | log.add_blobs([x], name='prelu_blob') 286 | layer = caffe_net.Layer_param(name=name, type='PReLU', 287 | bottom=bottom_blobs, top=[log.blobs(x)]) 288 | if weight.size()[0]==1: 289 | layer.param.prelu_param.channel_shared=True 290 | layer.add_data(weight.cpu().data.numpy()[0]) 291 | else: 292 | layer.add_data(weight.cpu().data.numpy()) 293 | log.cnet.add_layer(layer) 294 | return x 295 | 296 | def _leaky_relu(raw, input, negative_slope=0.01, inplace=False): 297 | x = raw(input, negative_slope) 298 | name = log.add_layer(name='leaky_relu') 299 | log.add_blobs([x], name='leaky_relu_blob') 300 | layer = caffe_net.Layer_param(name=name, type='ReLU', 301 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 302 | layer.param.relu_param.negative_slope=negative_slope 303 | log.cnet.add_layer(layer) 304 | return x 305 | 306 | def _tanh(raw, input): 307 | # for tanh activation 308 | x = raw(input) 309 | name = log.add_layer(name='tanh') 310 | log.add_blobs([x], name='tanh_blob') 311 | layer = caffe_net.Layer_param(name=name, type='TanH', 312 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 313 | log.cnet.add_layer(layer) 314 | return x 315 | 316 | def _softmax(raw, input, dim=None, _stacklevel=3): 317 | # for F.softmax 318 | x=raw(input, dim=dim) 319 | if dim is None: 320 | dim=F._get_softmax_dim('softmax', input.dim(), _stacklevel) 321 | bottom_blobs=[log.blobs(input)] 322 | name = log.add_layer(name='softmax') 323 | log.add_blobs([x], name='softmax_blob') 324 | layer = caffe_net.Layer_param(name=name, type='Softmax', 325 | bottom=bottom_blobs, top=[log.blobs(x)]) 326 | layer.param.softmax_param.axis=dim 327 | log.cnet.add_layer(layer) 328 | return x 329 | 330 | def _batch_norm(raw,input, running_mean, running_var, weight=None, bias=None, 331 | training=False, momentum=0.1, eps=1e-5): 332 | # because the runing_mean and runing_var will be changed after the _batch_norm operation, we first save the parameters 333 | 334 | x = raw(input, running_mean, running_var, weight, bias, 335 | training, momentum, eps) 336 | bottom_blobs = [log.blobs(input)] 337 | layer_name1 = log.add_layer(name='batch_norm') 338 | top_blobs = log.add_blobs([x], name='batch_norm_blob') 339 | layer1 = caffe_net.Layer_param(name=layer_name1, type='BatchNorm', 340 | bottom=bottom_blobs, top=top_blobs) 341 | if running_mean is None or running_var is None: 342 | # not use global_stats, normalization is performed over the current mini-batch 343 | layer1.batch_norm_param(use_global_stats=0,eps=eps) 344 | else: 345 | layer1.batch_norm_param(use_global_stats=1, eps=eps) 346 | running_mean_clone = running_mean.clone() 347 | running_var_clone = running_var.clone() 348 | layer1.add_data(running_mean_clone.cpu().numpy(), running_var_clone.cpu().numpy(), np.array([1.0])) 349 | log.cnet.add_layer(layer1) 350 | if weight is not None and bias is not None: 351 | layer_name2 = log.add_layer(name='bn_scale') 352 | layer2 = caffe_net.Layer_param(name=layer_name2, type='Scale', 353 | bottom=top_blobs, top=top_blobs) 354 | layer2.param.scale_param.bias_term = True 355 | layer2.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy()) 356 | log.cnet.add_layer(layer2) 357 | return x 358 | 359 | def _instance_norm(raw, input, running_mean=None, running_var=None, weight=None, 360 | bias=None, use_input_stats=True, momentum=0.1, eps=1e-5): 361 | # TODO: the batch size!=1 view operations 362 | print("WARNING: The Instance Normalization transfers to Caffe using BatchNorm, so the batch size should be 1") 363 | if running_var is not None or weight is not None: 364 | # TODO: the affine=True or track_running_stats=True case 365 | raise NotImplementedError("not implement the affine=True or track_running_stats=True case InstanceNorm") 366 | x= torch.batch_norm( 367 | input, weight, bias, running_mean, running_var, 368 | use_input_stats, momentum, eps,torch.backends.cudnn.enabled) 369 | bottom_blobs = [log.blobs(input)] 370 | layer_name1 = log.add_layer(name='instance_norm') 371 | top_blobs = log.add_blobs([x], name='instance_norm_blob') 372 | layer1 = caffe_net.Layer_param(name=layer_name1, type='BatchNorm', 373 | bottom=bottom_blobs, top=top_blobs) 374 | if running_mean is None or running_var is None: 375 | # not use global_stats, normalization is performed over the current mini-batch 376 | layer1.batch_norm_param(use_global_stats=0,eps=eps) 377 | running_mean=torch.zeros(input.size()[1]) 378 | running_var=torch.ones(input.size()[1]) 379 | else: 380 | layer1.batch_norm_param(use_global_stats=1, eps=eps) 381 | running_mean_clone = running_mean.clone() 382 | running_var_clone = running_var.clone() 383 | layer1.add_data(running_mean_clone.cpu().numpy(), running_var_clone.cpu().numpy(), np.array([1.0])) 384 | log.cnet.add_layer(layer1) 385 | if weight is not None and bias is not None: 386 | layer_name2 = log.add_layer(name='bn_scale') 387 | layer2 = caffe_net.Layer_param(name=layer_name2, type='Scale', 388 | bottom=top_blobs, top=top_blobs) 389 | layer2.param.scale_param.bias_term = True 390 | layer2.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy()) 391 | log.cnet.add_layer(layer2) 392 | return x 393 | 394 | 395 | #upsample layer 396 | def _interpolate(raw, input,size=None, scale_factor=None, mode='nearest', align_corners=None): 397 | # 定义的参数包括 scale,即输出与输入的尺寸比例,如 2;scale_h、scale_w, 398 | # 同 scale,分别为 h、w 方向上的尺寸比例;pad_out_h、pad_out_w,仅在 scale 为 2 时 399 | # 有用,对输出进行额外 padding 在 h、w 方向上的数值;upsample_h、upsample_w,输 400 | # 出图像尺寸的数值。在 Upsample 的相关代码中,推荐仅仅使用 upsample_h、 401 | # upsample_w 准确定义 Upsample 层的输出尺寸,其他所有的参数都不推荐继续使用。 402 | # for nearest _interpolate 403 | if mode != "nearest" or align_corners != None: 404 | raise NotImplementedError("not implement F.interpolate totoaly") 405 | x = raw(input,size , scale_factor ,mode) 406 | 407 | layer_name = log.add_layer(name='upsample') 408 | top_blobs = log.add_blobs([x], name='upsample_blob'.format(type)) 409 | layer = caffe_net.Layer_param(name=layer_name, type='Upsample', 410 | bottom=[log.blobs(input)], top=top_blobs) 411 | 412 | layer.upsample_param(size =(input.size(2),input.size(3)), scale_factor= scale_factor) 413 | log.cnet.add_layer(layer) 414 | return x 415 | 416 | 417 | #sigmid layer 418 | def _sigmoid(raw, input): 419 | # Applies the element-wise function: 420 | # 421 | # Sigmoid(x)= 1/(1+exp(−x)） 422 | # 423 | # ​ 424 | x = raw(input) 425 | name = log.add_layer(name='sigmoid') 426 | log.add_blobs([x], name='sigmoid_blob') 427 | layer = caffe_net.Layer_param(name=name, type='Sigmoid', 428 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 429 | log.cnet.add_layer(layer) 430 | return x 431 | 432 | #tanh layer 433 | def _tanh(raw, input): 434 | # Applies the element-wise function: 435 | # 436 | # torch.nn.Tanh 437 | # 438 | # ​ 439 | x = raw(input) 440 | name = log.add_layer(name='tanh') 441 | log.add_blobs([x], name='tanh_blob') 442 | layer = caffe_net.Layer_param(name=name, type='TanH', 443 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 444 | log.cnet.add_layer(layer) 445 | return x 446 | 447 | def _hardtanh(raw, input, min_val, max_val, inplace): 448 | # Applies the element-wise function: 449 | # 450 | # torch.nn.ReLu6 451 | # 452 | # ​ 453 | print('relu6: ', log.blobs(input)) 454 | x = raw(input, min_val, max_val) 455 | name = log.add_layer(name='relu6') 456 | log.add_blobs([x], name='relu6_blob') 457 | layer = caffe_net.Layer_param(name=name, type='ReLU6', 458 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 459 | log.cnet.add_layer(layer) 460 | return x 461 | 462 | #L2Norm layer 463 | def _l2Norm(raw, input, weight, eps): 464 | # Applies the element-wise function: 465 | # 466 | # L2Norm in vgg_ssd 467 | # 468 | # ​ 469 | x = raw(input, weight, eps) 470 | name = log.add_layer(name='normalize') 471 | log.add_blobs([x], name='normalize_blob') 472 | layer = caffe_net.Layer_param(name=name, type='Normalize', 473 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 474 | layer.norm_param(eps) 475 | 476 | layer.add_data(weight.cpu().data.numpy()) 477 | log.cnet.add_layer(layer) 478 | return x 479 | 480 | def _div(raw,inputs, inputs2): 481 | x=raw(inputs, inputs2) 482 | log.add_blobs([x],name='div_blob') 483 | return x 484 | 485 | 486 | # ----- for Variable operations -------- 487 | 488 | def _view(input, *args): 489 | x=raw_view(input, *args) 490 | if not NET_INITTED: 491 | return x 492 | layer_name=log.add_layer(name='view') 493 | top_blobs=log.add_blobs([x],name='view_blob') 494 | layer=caffe_net.Layer_param(name=layer_name,type='Reshape', 495 | bottom=[log.blobs(input)],top=top_blobs) 496 | # TODO: reshpae added to nn_tools layer 497 | dims=list(args) 498 | dims[0]=0 # the first dim should be batch_size 499 | layer.param.reshape_param.shape.CopyFrom(caffe_net.pb.BlobShape(dim=dims)) 500 | log.cnet.add_layer(layer) 501 | return x 502 | 503 | def _mean(input, *args,**kwargs): 504 | x=raw_mean(input, *args,**kwargs) 505 | if not NET_INITTED: 506 | return x 507 | layer_name=log.add_layer(name='mean') 508 | top_blobs=log.add_blobs([x],name='mean_blob') 509 | layer=caffe_net.Layer_param(name=layer_name,type='Reduction', 510 | bottom=[log.blobs(input)],top=top_blobs) 511 | if len(args)==1: 512 | dim=args[0] 513 | elif 'dim' in kwargs: 514 | dim=kwargs['dim'] 515 | else: 516 | raise NotImplementedError('mean operation must specify a dim') 517 | layer.param.reduction_param.operation=4 518 | layer.param.reduction_param.axis=dim 519 | log.cnet.add_layer(layer) 520 | return x 521 | 522 | def _add(input, *args): 523 | x = raw__add__(input, *args) 524 | if not NET_INITTED: 525 | return x 526 | layer_name = log.add_layer(name='add') 527 | top_blobs = log.add_blobs([x], name='add_blob') 528 | if log.blobs(args[0]) == None: 529 | log.add_blobs([args[0]], name='extra_blob') 530 | else: 531 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 532 | bottom=[log.blobs(input),log.blobs(args[0])], top=top_blobs) 533 | layer.param.eltwise_param.operation = 1 # sum is 1 534 | log.cnet.add_layer(layer) 535 | return x 536 | 537 | def _iadd(input, *args): 538 | x = raw__iadd__(input, *args) 539 | if not NET_INITTED: 540 | return x 541 | x=x.clone() 542 | layer_name = log.add_layer(name='add') 543 | top_blobs = log.add_blobs([x], name='add_blob') 544 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 545 | bottom=[log.blobs(input),log.blobs(args[0])], top=top_blobs) 546 | layer.param.eltwise_param.operation = 1 # sum is 1 547 | log.cnet.add_layer(layer) 548 | return x 549 | 550 | def _sub(input, *args): 551 | x = raw__sub__(input, *args) 552 | if not NET_INITTED: 553 | return x 554 | layer_name = log.add_layer(name='sub') 555 | top_blobs = log.add_blobs([x], name='sub_blob') 556 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 557 | bottom=[log.blobs(input),log.blobs(args[0])], top=top_blobs) 558 | layer.param.eltwise_param.operation = 1 # sum is 1 559 | layer.param.eltwise_param.coeff.extend([1.,-1.]) 560 | log.cnet.add_layer(layer) 561 | return x 562 | 563 | def _isub(input, *args): 564 | x = raw__isub__(input, *args) 565 | if not NET_INITTED: 566 | return x 567 | x=x.clone() 568 | layer_name = log.add_layer(name='sub') 569 | top_blobs = log.add_blobs([x], name='sub_blob') 570 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 571 | bottom=[log.blobs(input),log.blobs(args[0])], top=top_blobs) 572 | layer.param.eltwise_param.operation = 1 # sum is 1 573 | log.cnet.add_layer(layer) 574 | return x 575 | 576 | def _mul(input, *args): 577 | x = raw__mul__(input, *args) 578 | if not NET_INITTED: 579 | return x 580 | layer_name = log.add_layer(name='mul') 581 | top_blobs = log.add_blobs([x], name='mul_blob') 582 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 583 | bottom=[log.blobs(input), log.blobs(args[0])], top=top_blobs) 584 | layer.param.eltwise_param.operation = 0 # product is 1 585 | log.cnet.add_layer(layer) 586 | return x 587 | 588 | def _imul(input, *args): 589 | x = raw__imul__(input, *args) 590 | if not NET_INITTED: 591 | return x 592 | x = x.clone() 593 | layer_name = log.add_layer(name='mul') 594 | top_blobs = log.add_blobs([x], name='mul_blob') 595 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 596 | bottom=[log.blobs(input), log.blobs(args[0])], top=top_blobs) 597 | layer.param.eltwise_param.operation = 0 # product is 1 598 | layer.param.eltwise_param.coeff.extend([1., -1.]) 599 | log.cnet.add_layer(layer) 600 | return x 601 | 602 | 603 | #Permute layer 604 | def _permute(input, *args): 605 | x = raw__permute__(input, *args) 606 | name = log.add_layer(name='permute') 607 | log.add_blobs([x], name='permute_blob') 608 | layer = caffe_net.Layer_param(name=name, type='Permute', 609 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 610 | order1 = args[0] 611 | order2 = args[1] 612 | order3 = args[2] 613 | order4 = args[3] 614 | 615 | layer.permute_param(order1, order2, order3, order4) 616 | log.cnet.add_layer(layer) 617 | return x 618 | 619 | #contiguous 620 | def _contiguous(input, *args): 621 | x = raw__contiguous__(input, *args) 622 | name = log.add_layer(name='contiguous') 623 | log.add_blobs([x], name='contiguous_blob') 624 | layer = caffe_net.Layer_param(name=name, type='NeedRemove', 625 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 626 | log.cnet.add_layer(layer) 627 | return x 628 | 629 | #pow 630 | def _pow(input, *args): 631 | x = raw__pow__(input, *args) 632 | log.add_blobs([x], name='pow_blob') 633 | return x 634 | 635 | #sum 636 | def _sum(input, *args): 637 | x = raw__sum__(input, *args) 638 | log.add_blobs([x], name='sum_blob') 639 | return x 640 | 641 | # sqrt 642 | def _sqrt(input, *args): 643 | x = raw__sqrt__(input, *args) 644 | log.add_blobs([x], name='sqrt_blob') 645 | return x 646 | 647 | # unsqueeze 648 | def _unsqueeze(input, *args): 649 | x = raw__unsqueeze__(input, *args) 650 | log.add_blobs([x], name='unsqueeze_blob') 651 | return x 652 | 653 | def _expand_as(input, *args): 654 | # only support expand A(1, 1, H, W) to B(1, C, H, W) 655 | 656 | x = raw__expand_as__(input, *args) 657 | layer_name = log.add_layer(name="expand_as", with_num=True) 658 | log.add_blobs([x], name='expand_as_blob') 659 | layer = caffe_net.Layer_param(name=layer_name, type='Convolution', 660 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 661 | 662 | def constant_weight(shape): 663 | weights = np.ones(shape, dtype='float32') 664 | return weights 665 | 666 | channels = args[0].size(1) 667 | weight = constant_weight([channels, 1, 1, 1]) 668 | layer.conv_param(channels, kernel_size = 1, bias_term=False, weight_filler_type='xavier') 669 | layer.add_data(weight) 670 | log.cnet.add_layer(layer) 671 | return x 672 | 673 | # 核心组件，通过该类，实现对torch的function中的operators的输入，输出以及参数的读取 674 | class Rp(object): 675 | def __init__(self,raw,replace,**kwargs): 676 | # replace the raw function to replace function 677 | self.obj=replace 678 | self.raw=raw 679 | 680 | def __call__(self,*args,**kwargs): 681 | if not NET_INITTED: 682 | return self.raw(*args,**kwargs) 683 | for stack in traceback.walk_stack(None): 684 | if 'self' in stack[0].f_locals: 685 | layer=stack[0].f_locals['self'] 686 | if layer in layer_names: 687 | log.pytorch_layer_name=layer_names[layer] 688 | print(layer_names[layer]) 689 | break 690 | out=self.obj(self.raw,*args,**kwargs) 691 | # if isinstance(out,Variable): 692 | # out=[out] 693 | return out 694 | 695 | 696 | F.conv2d=Rp(F.conv2d,_conv2d) 697 | F.linear=Rp(F.linear,_linear) 698 | F.relu=Rp(F.relu,_relu) 699 | F.leaky_relu=Rp(F.leaky_relu,_leaky_relu) 700 | F.max_pool2d=Rp(F.max_pool2d,_max_pool2d) 701 | F.avg_pool2d=Rp(F.avg_pool2d,_avg_pool2d) 702 | F.adaptive_avg_pool2d = Rp(F.adaptive_avg_pool2d,_adaptive_avg_pool2d) 703 | F.dropout=Rp(F.dropout,_dropout) 704 | F.threshold=Rp(F.threshold,_threshold) 705 | F.prelu=Rp(F.prelu,_prelu) 706 | F.batch_norm=Rp(F.batch_norm,_batch_norm) 707 | F.instance_norm=Rp(F.instance_norm,_instance_norm) 708 | F.softmax=Rp(F.softmax,_softmax) 709 | F.conv_transpose2d=Rp(F.conv_transpose2d,_conv_transpose2d) 710 | F.interpolate = Rp(F.interpolate,_interpolate) 711 | F.sigmoid = Rp(F.sigmoid,_sigmoid) 712 | F.tanh = Rp(F.tanh,_tanh) 713 | F.tanh = Rp(F.tanh,_tanh) 714 | F.hardtanh = Rp(F.hardtanh,_hardtanh) 715 | # F.l2norm = Rp(F.l2norm,_l2Norm) 716 | 717 | torch.split=Rp(torch.split,_split) 718 | torch.max=Rp(torch.max,_max) 719 | torch.cat=Rp(torch.cat,_cat) 720 | torch.div=Rp(torch.div,_div) 721 | 722 | # TODO: other types of the view function 723 | try: 724 | raw_view=Variable.view 725 | Variable.view=_view 726 | raw_mean=Variable.mean 727 | Variable.mean=_mean 728 | raw__add__=Variable.__add__ 729 | Variable.__add__=_add 730 | raw__iadd__=Variable.__iadd__ 731 | Variable.__iadd__=_iadd 732 | raw__sub__=Variable.__sub__ 733 | Variable.__sub__=_sub 734 | raw__isub__=Variable.__isub__ 735 | Variable.__isub__=_isub 736 | raw__mul__ = Variable.__mul__ 737 | Variable.__mul__ = _mul 738 | raw__imul__ = Variable.__imul__ 739 | Variable.__imul__ = _imul 740 | except: 741 | # for new version 0.4.0 and later version 742 | for t in [torch.Tensor]: 743 | raw_view = t.view 744 | t.view = _view 745 | raw_mean = t.mean 746 | t.mean = _mean 747 | raw__add__ = t.__add__ 748 | t.__add__ = _add 749 | raw__iadd__ = t.__iadd__ 750 | t.__iadd__ = _iadd 751 | raw__sub__ = t.__sub__ 752 | t.__sub__ = _sub 753 | raw__isub__ = t.__isub__ 754 | t.__isub__ = _isub 755 | raw__mul__ = t.__mul__ 756 | t.__mul__=_mul 757 | raw__imul__ = t.__imul__ 758 | t.__imul__ = _imul 759 | raw__permute__ = t.permute 760 | t.permute = _permute 761 | raw__contiguous__ = t.contiguous 762 | t.contiguous = _contiguous 763 | raw__pow__ = t.pow 764 | t.pow = _pow 765 | raw__sum__ = t.sum 766 | t.sum = _sum 767 | raw__sqrt__ = t.sqrt 768 | t.sqrt = _sqrt 769 | raw__unsqueeze__ = t.unsqueeze 770 | t.unsqueeze = _unsqueeze 771 | raw__expand_as__ = t.expand_as 772 | t.expand_as = _expand_as 773 | 774 | 775 | def trans_net(net,input_var,name='TransferedPytorchModel'): 776 | print('Starting Transform, This will take a while') 777 | log.init([input_var]) 778 | log.cnet.net.name=name 779 | log.cnet.net.input.extend([log.blobs(input_var)]) 780 | log.cnet.net.input_dim.extend(input_var.size()) 781 | global NET_INITTED 782 | NET_INITTED=True 783 | for name,layer in net.named_modules(): 784 | layer_names[layer]=name 785 | print("torch ops name:", layer_names) 786 | out = net.forward(input_var) 787 | print('Transform Completed') 788 | 789 | def save_prototxt(save_name): 790 | log.cnet.remove_layer_by_type("NeedRemove") 791 | log.cnet.save_prototxt(save_name) 792 | 793 | def save_caffemodel(save_name): 794 | log.cnet.save(save_name) 795 | --------------------------------------------------------------------------------