├── .gitignore ├── Caffe ├── ReadMe.md ├── __init__.py ├── caffe.proto ├── caffe_lmdb.py ├── caffe_net.py ├── caffe_pb2.py ├── layer_param.py └── net.py ├── Datasets ├── README.md ├── __init__.py ├── adience.py ├── base_datasets.py ├── cifar10.py ├── imagenet.py ├── imdb_wiki.py ├── lmdb_data.proto ├── lmdb_data_pb2.py ├── lmdb_datasets.py ├── mnist.py └── wider.py ├── Pytorch ├── __init__.py ├── augmentations.py ├── eval.py ├── train.py └── utils.py ├── README.md ├── Tensorflow ├── __init__.py ├── constructor.py └── graph.py ├── __init__.py ├── analysis ├── CaffeA.py ├── MxnetA.py ├── PytorchA.py ├── README.md ├── __init__.py ├── blob.py ├── layers.py ├── roi.py └── utils.py ├── caffe_analyser.py ├── example ├── MGN-annalyze.py ├── MGN_2_caffe.py ├── MGN_analysis_example.py ├── alexnet_pytorch_to_caffe.py ├── inceptionv3_pytorch_to_caffe.py ├── mgn.net ├── option.py ├── resnet_pytorch_2_caffe.py ├── resnet_pytorch_analysis_example.py ├── verify_deploy.py └── vgg19_pytorch_to_caffe.py ├── funcs.py ├── keras_to_caffe.py ├── model ├── __init__.py ├── mgn.py └── resnet.py ├── mxnet_analyser.py ├── pytorch_analyse.csv ├── pytorch_analyser.py ├── pytorch_to_caffe.py └── utils ├── __init__.py ├── functions.py ├── random_erasing.py ├── re_ranking.py └── utility.py /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | *.pyc 3 | output/* 4 | *.pt 5 | *.kmodel 6 | *.caffemodel 7 | *.prototxt 8 | tmp/* 9 | nohup.out -------------------------------------------------------------------------------- /Caffe/ReadMe.md: -------------------------------------------------------------------------------- 1 | # The Caffe in nn_tools 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/wdd0225/pytorch2caffe/479a18308903e502f74dce6d2b4c8dede2ef62ee/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 | 11 | def layer_index(self,layer_name): 12 | # find a layer's index by name. if the layer was found, return the layer position in the net, else return -1. 13 | for i, layer in enumerate(self.net.layer): 14 | if layer.name == layer_name: 15 | return i 16 | 17 | def add_layer(self,layer_params,before='',after=''): 18 | # find the before of after layer's position 19 | index = -1 20 | if after != '': 21 | index = self.layer_index(after) + 1 22 | if before != '': 23 | index = self.layer_index(before) 24 | new_layer = pb.LayerParameter() 25 | new_layer.CopyFrom(layer_params.param) 26 | #insert the layer into the layer protolist 27 | if index != -1: 28 | self.net.layer.add() 29 | for i in range(len(self.net.layer) - 1, index, -1): 30 | self.net.layer[i].CopyFrom(self.net.layer[i - 1]) 31 | self.net.layer[index].CopyFrom(new_layer) 32 | else: 33 | self.net.layer.extend([new_layer]) 34 | 35 | def remove_layer_by_name(self,layer_name): 36 | for i,layer in enumerate(self.net.layer): 37 | if layer.name == layer_name: 38 | del self.net.layer[i] 39 | return 40 | raise(AttributeError, "cannot found layer %s" % str(layer_name)) 41 | 42 | def get_layer_by_name(self, layer_name): 43 | # get the layer by layer_name 44 | for layer in self.net.layer: 45 | if layer.name == layer_name: 46 | return layer 47 | raise(AttributeError, "cannot found layer %s" % str(layer_name)) 48 | 49 | def save_prototxt(self,path): 50 | prototxt=pb.NetParameter() 51 | prototxt.CopyFrom(self.net) 52 | for layer in prototxt.layer: 53 | del layer.blobs[:] 54 | with open(path,'w') as f: 55 | f.write(text_format.MessageToString(prototxt)) 56 | 57 | def layer(self,layer_name): 58 | return self.get_layer_by_name(layer_name) 59 | 60 | def layers(self): 61 | return list(self.net.layer) 62 | 63 | 64 | 65 | class Prototxt(_Net): 66 | def __init__(self,file_name=''): 67 | super(Prototxt,self).__init__() 68 | self.file_name=file_name 69 | if file_name!='': 70 | f = open(file_name,'r') 71 | text_format.Parse(f.read(), self.net) 72 | pass 73 | 74 | def init_caffemodel(self,caffe_cmd_path='caffe'): 75 | """ 76 | :param caffe_cmd_path: The shell command of caffe, normally at /build/tools/caffe 77 | """ 78 | s=pb.SolverParameter() 79 | s.train_net=self.file_name 80 | s.max_iter=0 81 | s.base_lr=1 82 | s.solver_mode = pb.SolverParameter.CPU 83 | s.snapshot_prefix='./nn' 84 | with open('/tmp/nn_tools_solver.prototxt','w') as f: 85 | f.write(str(s)) 86 | import os 87 | os.system('%s train --solver /tmp/nn_tools_solver.prototxt'%caffe_cmd_path) 88 | 89 | class Caffemodel(_Net): 90 | def __init__(self, file_name=''): 91 | super(Caffemodel,self).__init__() 92 | # caffe_model dir 93 | if file_name!='': 94 | f = open(file_name,'rb') 95 | self.net.ParseFromString(f.read()) 96 | f.close() 97 | 98 | def save(self, path): 99 | with open(path,'wb') as f: 100 | f.write(self.net.SerializeToString()) 101 | 102 | def add_layer_with_data(self,layer_params,datas, before='', after=''): 103 | """ 104 | Args: 105 | layer_params:A Layer_Param object 106 | datas:a fixed dimension numpy object list 107 | after: put the layer after a specified layer 108 | before: put the layer before a specified layer 109 | """ 110 | self.add_layer(layer_params,before,after) 111 | new_layer =self.layer(layer_params.name) 112 | 113 | #process blobs 114 | del new_layer.blobs[:] 115 | for data in datas: 116 | new_blob=new_layer.blobs.add() 117 | for dim in data.shape: 118 | new_blob.shape.dim.append(dim) 119 | new_blob.data.extend(data.flatten().astype(float)) 120 | 121 | def get_layer_data(self,layer_name): 122 | layer=self.layer(layer_name) 123 | datas=[] 124 | for blob in layer.blobs: 125 | shape=list(blob.shape.dim) 126 | data=np.array(blob.data).reshape(shape) 127 | datas.append(data) 128 | return datas 129 | 130 | def set_layer_data(self,layer_name,datas): 131 | # datas is normally a list of [weights,bias] 132 | layer=self.layer(layer_name) 133 | for blob,data in zip(layer.blobs,datas): 134 | blob.data[:]=data.flatten() 135 | pass 136 | 137 | class Net(): 138 | def __init__(self,*args,**kwargs): 139 | raise(TypeError,'the class Net is no longer used, please use Caffemodel or Prototxt instead') -------------------------------------------------------------------------------- /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=False): 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 {} is not the same,try hieht and wight spilt up".format(item)) 13 | return item 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'): 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_filler.type = bias_filler 42 | self.param.inner_product_param.CopyFrom(fc_param) 43 | 44 | def conv_param(self, num_output, kernel_size, stride=(1), pad=(0,), 45 | weight_filler_type='xavier', bias_filler_type='constant', 46 | bias_term=True, dilation=None,groups=None): 47 | """ 48 | add a conv_param layer if you spec the layer type "Convolution" 49 | Args: 50 | num_output: a int 51 | kernel_size: int list 52 | stride: a int list 53 | weight_filler_type: the weight filer type 54 | bias_filler_type: the bias filler type 55 | Returns: 56 | """ 57 | if self.type!='Convolution': 58 | raise TypeError('the layer type must be Convolution if you want set conv param') 59 | conv_param=pb.ConvolutionParameter() 60 | conv_param.num_output=num_output 61 | conv_param.kernel_size.extend(pair_reduce(kernel_size)) 62 | conv_param.stride.extend(pair_reduce(stride)) 63 | conv_param.pad.extend(pair_reduce(pad)) 64 | conv_param.bias_term=bias_term 65 | conv_param.weight_filler.type=weight_filler_type 66 | if bias_term: 67 | conv_param.bias_filler.type = bias_filler_type 68 | if dilation: 69 | conv_param.dilation.extend(pair_reduce(dilation)) 70 | if groups: 71 | conv_param.group=groups 72 | self.param.convolution_param.CopyFrom(conv_param) 73 | 74 | def pool_param(self,type='MAX',kernel_size=2,stride=2,pad=None): 75 | pool_param=pb.PoolingParameter() 76 | pool_param.pool=pool_param.PoolMethod.Value(type) 77 | if len(pair_process(kernel_size)) > 1: 78 | pool_param.kernel_h = kernel_size[0] 79 | pool_param.kernel_w = kernel_size[1] 80 | else: 81 | pool_param.kernel_size=kernel_size 82 | if len(pair_process(stride)) > 1: 83 | pool_param.stride_h = stride[0] 84 | pool_param.stride_w = stride[1] 85 | else: 86 | pool_param.stride=stride 87 | 88 | if pad: 89 | pool_param.pad=pad 90 | self.param.pooling_param.CopyFrom(pool_param) 91 | 92 | def batch_norm_param(self,use_global_stats=0,moving_average_fraction=None,eps=None): 93 | bn_param=pb.BatchNormParameter() 94 | bn_param.use_global_stats=use_global_stats 95 | if moving_average_fraction: 96 | bn_param.moving_average_fraction=moving_average_fraction 97 | if eps: 98 | bn_param.eps = eps 99 | self.param.batch_norm_param.CopyFrom(bn_param) 100 | 101 | def add_data(self,*args): 102 | """Args are data numpy array 103 | """ 104 | del self.param.blobs[:] 105 | for data in args: 106 | new_blob = self.param.blobs.add() 107 | for dim in data.shape: 108 | new_blob.shape.dim.append(dim) 109 | new_blob.data.extend(data.flatten().astype(float)) 110 | 111 | def set_params_by_dict(self,dic): 112 | pass 113 | 114 | def copy_from(self,layer_param): 115 | pass 116 | 117 | def set_enum(param,key,value): 118 | setattr(param,key,param.Value(value)) 119 | -------------------------------------------------------------------------------- /Caffe/net.py: -------------------------------------------------------------------------------- 1 | raise ImportError,'the nn_tools.Caffe.net is no longer used, please use nn_tools.Caffe.caffe_net' -------------------------------------------------------------------------------- /Datasets/README.md: -------------------------------------------------------------------------------- 1 | == Datasets API == 2 | 3 | Providing some API for specified datasets. 4 | 5 | === IMDB-WIKI === 6 | 7 | API in `imdb_wiki.py`. 8 | 9 | - read_mat: Reading the .mat annotation files in IMDB-WIKI datasets and converting them to python objects. 10 | Then store in the cache file (default in `/tmp/imdb_wiki.pth`). It will also return the objects. 11 | The format of the data is `[full_paths_list, face_locations_list, genders_list, ages_list]` 12 | -------------------------------------------------------------------------------- /Datasets/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wdd0225/pytorch2caffe/479a18308903e502f74dce6d2b4c8dede2ef62ee/Datasets/__init__.py -------------------------------------------------------------------------------- /Datasets/adience.py: -------------------------------------------------------------------------------- 1 | import scipy.io as sio 2 | import os 3 | import cPickle as pickle 4 | import numpy as np 5 | 6 | AGES_START=[0,4,8,15,25,38,48,60] # the start of the age scope 7 | AGES_END=[2,6,13,20,32,43,53,100] # the end of the age scope 8 | AGES_MID=np.array([(i+j)/2 for i,j in zip(AGES_END,AGES_START)]) 9 | GENDER=['m','f'] # u for unknown 10 | 11 | def get_class(age): 12 | for i,(s,e) in enumerate(zip(AGES_START,AGES_END)): 13 | if s<=age<=e: 14 | return i 15 | return np.argmin(np.abs(AGES_MID-age)) 16 | 17 | def read_txt(root_path,cache_dir='/tmp'): 18 | # Reading the .txt annotation files in adience datasets and converting them to python objects. 19 | # Storing in the cache file (default in `/tmp/adience.pth`). 20 | txt_list=['fold_0_data.txt','fold_1_data.txt','fold_2_data.txt','fold_3_data.txt','fold_4_data.txt'] 21 | cache_file = os.path.join(cache_dir, 'adience.pth') 22 | 23 | if not os.path.isfile(cache_file): 24 | print "generating cache_file" 25 | image_paths, face_ids, ages, genders, locations, angs, yaws, scores=[],[],[],[],[],[],[],[] 26 | for txt in txt_list: 27 | txt_path = os.path.join(root_path, txt) 28 | with open(txt_path) as f: 29 | lines = f.readlines()[1:] 30 | for line in lines: 31 | user_id, original_image, face_id, age, gender, x,\ 32 | y, dx, dy, tilt_ang, yaw, fiducial_score = line.split('\t') 33 | age=eval(age) 34 | if type(age)==tuple: 35 | try:age=AGES_START.index(age[0]) 36 | except: 37 | continue 38 | elif type(age)==int: 39 | age=np.argmin(np.abs(AGES_MID-age)) 40 | else: 41 | continue # include None 42 | try: gender=GENDER.index(gender) 43 | except: 44 | continue 45 | image_paths.append(os.path.join(user_id, original_image)) 46 | face_ids.append(int(face_id)) 47 | ages.append(age) 48 | genders.append(gender) 49 | locations.append((int(x),int(y),int(dx),int(dy))) 50 | angs.append(int(tilt_ang)) 51 | yaws.append(int(yaw)) 52 | scores.append(int(fiducial_score)) 53 | save_obj=[image_paths, face_ids, ages, genders, locations, angs, yaws, scores] 54 | pickle.dump(save_obj, open(cache_file, 'wb')) 55 | else: 56 | print "read from cache_file" 57 | image_paths, face_ids, ages, genders, locations, angs, yaws, scores = pickle.load(open(cache_file, 'rb')) 58 | print "read mat OK" 59 | return image_paths, face_ids, ages, genders, locations, angs, yaws, scores 60 | 61 | if __name__=='__main__': 62 | import argparse 63 | parser=argparse.ArgumentParser() 64 | parser.add_argument('root_path',help='root path of adience datasets',type=str) 65 | args=parser.parse_args() 66 | read_txt(args.root_path) 67 | -------------------------------------------------------------------------------- /Datasets/base_datasets.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | class BaseDataset(): 4 | # the father of all datasets class, there are some interface was defined 5 | def __init__(self): 6 | # the data saved in the class was a list with the format [data,label,...] 7 | self.data={} 8 | self.init_data() 9 | self.p={'train':0,'test':0,'val':0} 10 | 11 | def init_data(self): 12 | pass 13 | 14 | def get_test_data(self): 15 | pass 16 | 17 | def get_train_data(self): 18 | pass 19 | 20 | def get_val_data(self): 21 | pass 22 | 23 | def _shuffle(self, x): 24 | 25 | idx = np.arange(len(x[0])) 26 | np.random.shuffle(idx) 27 | for i, xi in enumerate(x): 28 | x[i] = xi[idx] 29 | return x 30 | 31 | def batch(self, name, batch_size): 32 | p = self.p[name] 33 | if p + batch_size >= len(self.data[name]): 34 | self.p[name] = 0 35 | self.data[name] = self._shuffle(self.data[name]) 36 | self.p[name] += batch_size 37 | return [i[p:p+batch_size] for i in self.data[name]] -------------------------------------------------------------------------------- /Datasets/cifar10.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | class Cifar10(): 4 | def __init__(self): 5 | pass -------------------------------------------------------------------------------- /Datasets/imagenet.py: -------------------------------------------------------------------------------- 1 | from . import lmdb_datasets 2 | from torchvision import datasets,transforms 3 | import os 4 | import os.path as osp 5 | import torch.utils.data 6 | import numpy as np 7 | import cv2 8 | from . import lmdb_data_pb2 as pb2 9 | import Queue 10 | import time 11 | import multiprocessing 12 | 13 | DATASET_SIZE=100 14 | mean=np.array([[[0.485]], [[0.456]], [[0.406]]])*255 15 | std=np.array([[[0.229]], [[0.224]], [[0.225]]])*255 16 | 17 | class Imagenet_LMDB(lmdb_datasets.LMDB): 18 | def __init__(self,imagenet_dir,train=False): 19 | self.train_name='imagenet_train_lmdb' 20 | self.val_name='imagenet_val_lmdb' 21 | self.train=train 22 | super(Imagenet_LMDB, self).__init__(osp.join(imagenet_dir,train and self.train_name or self.val_name)) 23 | txn=self.env.begin() 24 | self.cur=txn.cursor() 25 | self.data = Queue.Queue(DATASET_SIZE*2) 26 | self.target = Queue.Queue(DATASET_SIZE*2) 27 | self.point=0 28 | # self._read_from_lmdb() 29 | 30 | def data_transfrom(self,data,other): 31 | data=data.astype(np.float32) 32 | if self.train: 33 | shape=np.fromstring(other[0],np.uint16) 34 | data=data.reshape(shape) 35 | # Random crop 36 | _, w, h = data.shape 37 | x1 = np.random.randint(0, w - 224) 38 | y1 = np.random.randint(0, h - 224) 39 | data=data[:,x1:x1+224 ,y1:y1 + 224] 40 | # HorizontalFlip 41 | #TODO horizontal flip 42 | else: 43 | data = data.reshape([3, 224, 224]) 44 | data = (data - mean) / std 45 | tensor = torch.Tensor(data) 46 | del data 47 | return tensor 48 | 49 | def target_transfrom(self,target): 50 | return target 51 | 52 | def _read_from_lmdb(self): 53 | self.cur.next() 54 | if not self.cur.key(): 55 | self.cur.first() 56 | dataset = pb2.Dataset().FromString(self.cur.value()) 57 | for datum in dataset.datums: 58 | data = np.fromstring(datum.data, np.uint8) 59 | try: 60 | data = self.data_transfrom(data, datum.other) 61 | except: 62 | print 'cannot trans ', data.shape 63 | continue 64 | target = int(datum.target) 65 | target = self.target_transfrom(target) 66 | self.data.put(data) 67 | self.target.put(target) 68 | # print 'read_from_lmdb', time.time()-r 69 | del dataset 70 | 71 | # def read_from_lmdb(self): 72 | # process=multiprocessing.Process(target=self._read_from_lmdb) 73 | # process.start() 74 | 75 | def __getitem__(self,index): 76 | if self.data.qsize()0: 46 | loc[1]-=h*expand_rate 47 | loc[3]+=h*expand_rate 48 | loc[0]-=w*expand_rate 49 | loc[2]+=w*expand_rate 50 | loc=np.maximum(0,loc) 51 | loc[3]=np.minimum(im.shape[0],loc[3]) 52 | loc[2]=np.minimum(im.shape[1],loc[2]) 53 | # loc=loc.astype(np.int32) 54 | im=im[loc[1]:loc[3],loc[0]:loc[2]] 55 | h = loc[3] - loc[1] 56 | w = loc[2] - loc[0] 57 | if w>max_size or h>max_size: 58 | if w!=h: 59 | pass 60 | print("resize picture %s"%image_path) 61 | resize_factor=np.minimum(1.*max_size/w,1.*max_size/h) 62 | im=cv2.resize(im,(int(w*resize_factor),int(h*resize_factor))) 63 | cv2.imwrite(out_path,im) 64 | 65 | 66 | 67 | def generate_caffe_txt_age(mat_path,output_path,age_range, 68 | cache_dir='/tmp',ignore_second_face=False,test_ratio=0.2): 69 | # read mat file then generate train.txt and test.txt for age estimation training. 70 | # `age_range` is a list containing age range like (0,2) and (32,40). 71 | image_paths, face_locations, genders, ages, face_scores, second_face_scores = read_mat(mat_path,cache_dir) 72 | # generate classes 73 | ages_mid=np.array([(age[0]+age[1])/2 for age in age_range]) 74 | classes=[] 75 | for age in ages: 76 | classes.append(np.argmin(np.abs(ages_mid-age))) 77 | for idx,r in enumerate(age_range): 78 | if r[0]<=age<=r[1]: 79 | classes[-1]=idx 80 | break 81 | shuffle_idx=np.arange(len(image_paths)) 82 | np.random.shuffle(shuffle_idx) 83 | train_idx=shuffle_idx[int(len(shuffle_idx)*test_ratio):] 84 | test_idx=shuffle_idx[:int(len(shuffle_idx)*test_ratio)] 85 | with open(os.path.join(output_path, 'age_train.txt'), 'w') as trainf: 86 | for idx in train_idx: 87 | if ignore_second_face: 88 | if second_face_scores[idx]>1.5: 89 | continue 90 | trainf.write("%s %d\n" % (image_paths[idx], classes[idx])) 91 | with open(os.path.join(output_path, 'age_test.txt'), 'w') as testf: 92 | for idx in test_idx: 93 | if ignore_second_face: 94 | if second_face_scores[idx]>1.5: 95 | continue 96 | testf.write("%s %d\n"%(image_paths[idx],classes[idx])) 97 | 98 | def generate_caffe_txt_gender(mat_path,output_path,cache_dir='/tmp',test_ratio=0.2): 99 | # read mat file then generate train.txt and test.txt for gender estimation training 100 | image_paths, face_locations, genders, ages, face_score, second_face_score = read_mat(mat_path, cache_dir) 101 | shuffle_idx=np.arange(len(image_paths)) 102 | np.random.shuffle(shuffle_idx) 103 | train_idx=shuffle_idx[int(len(shuffle_idx)*test_ratio):] 104 | test_idx=shuffle_idx[:int(len(shuffle_idx)*test_ratio)] 105 | with open(os.path.join(output_path, 'gender_train.txt'), 'w') as trainf: 106 | for idx in train_idx: 107 | if genders[idx]==-1:continue 108 | trainf.write("%s %d\n" % (image_paths[idx], genders[idx])) 109 | with open(os.path.join(output_path, 'gender_test.txt'), 'w') as testf: 110 | for idx in test_idx: 111 | if genders[idx]==-1:continue 112 | testf.write("%s %d\n"%(image_paths[idx],genders[idx])) 113 | 114 | if __name__=='__main__': 115 | import argparse 116 | parser=argparse.ArgumentParser() 117 | parser.add_argument('mat_path',help='.mat file path in IMDB-WIKI datasets',type=str) 118 | args=parser.parse_args() 119 | read_mat(args.mat_path) 120 | -------------------------------------------------------------------------------- /Datasets/lmdb_data.proto: -------------------------------------------------------------------------------- 1 | // this file is for nn_tools lmdb_dataset 2 | // written by yuanzhihang1@126.com 3 | 4 | syntax = "proto2"; 5 | 6 | message Datum{ 7 | required bytes target=1; 8 | required bytes data=2; 9 | repeated bytes other=3; 10 | } 11 | 12 | message Dataset{ 13 | repeated Datum datums=1; 14 | } -------------------------------------------------------------------------------- /Datasets/lmdb_data_pb2.py: -------------------------------------------------------------------------------- 1 | # Generated by the protocol buffer compiler. DO NOT EDIT! 2 | # source: lmdb_data.proto 3 | 4 | import sys 5 | _b=sys.version_info[0]<3 and (lambda x:x) or (lambda x:x.encode('latin1')) 6 | from google.protobuf import descriptor as _descriptor 7 | from google.protobuf import message as _message 8 | from google.protobuf import reflection as _reflection 9 | from google.protobuf import symbol_database as _symbol_database 10 | from google.protobuf import descriptor_pb2 11 | # @@protoc_insertion_point(imports) 12 | 13 | _sym_db = _symbol_database.Default() 14 | 15 | 16 | 17 | 18 | DESCRIPTOR = _descriptor.FileDescriptor( 19 | name='lmdb_data.proto', 20 | package='', 21 | serialized_pb=_b('\n\x0flmdb_data.proto\"4\n\x05\x44\x61tum\x12\x0e\n\x06target\x18\x01 \x02(\x0c\x12\x0c\n\x04\x64\x61ta\x18\x02 \x02(\x0c\x12\r\n\x05other\x18\x03 \x03(\x0c\"!\n\x07\x44\x61taset\x12\x16\n\x06\x64\x61tums\x18\x01 \x03(\x0b\x32\x06.Datum') 22 | ) 23 | _sym_db.RegisterFileDescriptor(DESCRIPTOR) 24 | 25 | 26 | 27 | 28 | _DATUM = _descriptor.Descriptor( 29 | name='Datum', 30 | full_name='Datum', 31 | filename=None, 32 | file=DESCRIPTOR, 33 | containing_type=None, 34 | fields=[ 35 | _descriptor.FieldDescriptor( 36 | name='target', full_name='Datum.target', index=0, 37 | number=1, type=12, cpp_type=9, label=2, 38 | has_default_value=False, default_value=_b(""), 39 | message_type=None, enum_type=None, containing_type=None, 40 | is_extension=False, extension_scope=None, 41 | options=None), 42 | _descriptor.FieldDescriptor( 43 | name='data', full_name='Datum.data', index=1, 44 | number=2, type=12, cpp_type=9, label=2, 45 | has_default_value=False, default_value=_b(""), 46 | message_type=None, enum_type=None, containing_type=None, 47 | is_extension=False, extension_scope=None, 48 | options=None), 49 | _descriptor.FieldDescriptor( 50 | name='other', full_name='Datum.other', index=2, 51 | number=3, type=12, cpp_type=9, label=3, 52 | has_default_value=False, default_value=[], 53 | message_type=None, enum_type=None, containing_type=None, 54 | is_extension=False, extension_scope=None, 55 | options=None), 56 | ], 57 | extensions=[ 58 | ], 59 | nested_types=[], 60 | enum_types=[ 61 | ], 62 | options=None, 63 | is_extendable=False, 64 | extension_ranges=[], 65 | oneofs=[ 66 | ], 67 | serialized_start=19, 68 | serialized_end=71, 69 | ) 70 | 71 | 72 | _DATASET = _descriptor.Descriptor( 73 | name='Dataset', 74 | full_name='Dataset', 75 | filename=None, 76 | file=DESCRIPTOR, 77 | containing_type=None, 78 | fields=[ 79 | _descriptor.FieldDescriptor( 80 | name='datums', full_name='Dataset.datums', index=0, 81 | number=1, type=11, cpp_type=10, label=3, 82 | has_default_value=False, default_value=[], 83 | message_type=None, enum_type=None, containing_type=None, 84 | is_extension=False, extension_scope=None, 85 | options=None), 86 | ], 87 | extensions=[ 88 | ], 89 | nested_types=[], 90 | enum_types=[ 91 | ], 92 | options=None, 93 | is_extendable=False, 94 | extension_ranges=[], 95 | oneofs=[ 96 | ], 97 | serialized_start=73, 98 | serialized_end=106, 99 | ) 100 | 101 | _DATASET.fields_by_name['datums'].message_type = _DATUM 102 | DESCRIPTOR.message_types_by_name['Datum'] = _DATUM 103 | DESCRIPTOR.message_types_by_name['Dataset'] = _DATASET 104 | 105 | Datum = _reflection.GeneratedProtocolMessageType('Datum', (_message.Message,), dict( 106 | DESCRIPTOR = _DATUM, 107 | __module__ = 'lmdb_data_pb2' 108 | # @@protoc_insertion_point(class_scope:Datum) 109 | )) 110 | _sym_db.RegisterMessage(Datum) 111 | 112 | Dataset = _reflection.GeneratedProtocolMessageType('Dataset', (_message.Message,), dict( 113 | DESCRIPTOR = _DATASET, 114 | __module__ = 'lmdb_data_pb2' 115 | # @@protoc_insertion_point(class_scope:Dataset) 116 | )) 117 | _sym_db.RegisterMessage(Dataset) 118 | 119 | 120 | # @@protoc_insertion_point(module_scope) 121 | -------------------------------------------------------------------------------- /Datasets/lmdb_datasets.py: -------------------------------------------------------------------------------- 1 | import caffe_lmdb 2 | from . import lmdb_data_pb2 as pb2 3 | import numpy as np 4 | import multiprocessing 5 | import os 6 | 7 | 8 | class LMDB(object): 9 | def __init__(self,lmdb_dir): 10 | self.env=caffe_lmdb.Environment(lmdb_dir, map_size=int(1e12)) 11 | 12 | # -------------------------------------- 13 | # for LMDB writer 14 | 15 | 16 | 17 | # -------------------------------------- 18 | # for LMDB reader 19 | class LMDB_generator(object): 20 | def __init__(self,lmdb_dir): 21 | self.env=caffe_lmdb.Environment(lmdb_dir, map_size=int(1e12)) 22 | 23 | def generate_datum(self,data,target,other=None): 24 | datum = pb2.Datum() 25 | datum.data=data 26 | datum.target=target 27 | if other: 28 | datum.other.extend(other) 29 | # data=np.fromstring(datum.data,np.uint8).reshape(datum.other[0],np.uint16) 30 | return datum 31 | 32 | def generate_dataset(self,datas,targets,others=None): 33 | dataset = pb2.Dataset() 34 | assert len(datas)==len(targets),ValueError('the lengths of datas and targets are not the same') 35 | for idx in xrange(len(datas)): 36 | try: 37 | if others==None: 38 | datum=self.generate_datum(datas[idx],targets[idx]) 39 | else: 40 | datum = self.generate_datum(datas[idx], targets[idx],others[idx]) 41 | except: 42 | print('generate the datum failed at %d, continue it'%idx) 43 | continue 44 | dataset.datums.extend([datum]) 45 | return dataset 46 | 47 | def commit_dataset(self,dataset,idx): 48 | txn=self.env.begin(write=True) 49 | txn.put(str(idx),dataset.SerializeToString()) 50 | txn.commit() 51 | 52 | def write_classification_lmdb(self, data_loader, num_per_dataset=3000, write_shape=False): 53 | # torch_data_loader are iterator that iterates a (data,target) 54 | # data should be a numpy array 55 | # target should be a int number 56 | datas=[] 57 | targets=[] 58 | others=[] 59 | for idx,(data,target) in enumerate(data_loader): 60 | datas.append(data.tobytes()) 61 | targets.append(bytes(target)) 62 | if write_shape: 63 | others.append([np.array(data.shape, np.uint16).tobytes()]) 64 | if (idx%num_per_dataset==0 and idx!=0) or (idx==len(data_loader)-1): 65 | print('lmdb write at image %d'%idx) 66 | dataset=self.generate_dataset(datas,targets,write_shape and others or None) 67 | self.commit_dataset(dataset,np.ceil(1.*idx/num_per_dataset)) 68 | datas=[] 69 | targets=[] 70 | others = [] 71 | 72 | def write_lmdb_mutiprocess(self,torch_data_loader,num_thread=10,num_per_dataset=6000): 73 | # TODO mutiprocess write lmdb 74 | pass 75 | 76 | 77 | if __name__=='__main__': 78 | pass -------------------------------------------------------------------------------- /Datasets/mnist.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | import struct 4 | import base_datasets 5 | 6 | class Mnist(base_datasets.BaseDataset): 7 | # read the raw mnist data 8 | def __init__(self, dir, only_test=0): 9 | # dir should include t10k-images-idx3-ubyte t10k-labels-idx1-ubyte train-images-idx3-ubyte train-labels-idx1-ubyte 10 | # the data can be download at http://yann.lecun.com/exdb/mnist/ 11 | self.dir=dir 12 | self.train_dir=[os.path.join(dir,'train-images-idx3-ubyte'), 13 | os.path.join(dir,'train-labels-idx1-ubyte')] 14 | self.test_dir = [os.path.join(dir, 't10k-images-idx3-ubyte'), 15 | os.path.join(dir, 't10k-labels-idx1-ubyte')] 16 | base_datasets.BaseDataset.__init__(self) 17 | self.test_data=self.get_test_data() 18 | if not only_test: 19 | self.train_data=self.get_train_data() 20 | self.ptest=0 21 | self.ptrain=0 22 | 23 | def norm(self,x): 24 | data = x.reshape([-1, 28, 28, 1]).astype(np.float32) 25 | data = (data - 167) / 167 26 | return data 27 | 28 | def get_train_data(self,norm=1): 29 | #return the train dataset as a list [images,labels] 30 | data=self.read_images(self.train_dir[0]) 31 | if norm==1: 32 | data=self.norm(data) 33 | return [data,self.read_labels(self.train_dir[1])] 34 | 35 | def get_test_data(self,norm=1): 36 | #return the test dataset as a list [images,labels] 37 | data=self.read_images(self.test_dir[0]) 38 | if norm==1: 39 | data=self.norm(data) 40 | return [data,self.read_labels(self.test_dir[1])] 41 | 42 | def shuffle(self,x): 43 | idx=np.arange(len(x[0])) 44 | np.random.shuffle(idx) 45 | for i,xi in enumerate(x): 46 | x[i]=xi[idx] 47 | return x 48 | 49 | def _batch(self,data,p,batch_size): 50 | return [i[p-batch_size:p] for i in data] 51 | 52 | def get_test_batch(self,batch_size): 53 | if self.ptest+batch_size>=len(self.test_data): 54 | self.ptest=0 55 | self.test_data=self.shuffle(self.test_data) 56 | self.ptest+=batch_size 57 | return self._batch(self.test_data,self.ptest,batch_size) 58 | 59 | def get_train_batch(self, batch_size): 60 | if self.ptrain + batch_size >= len(self.train_data): 61 | self.ptrain = 0 62 | self.train_data=self.shuffle(self.train_data) 63 | self.ptrain += batch_size 64 | return self._batch(self.train_data,self.ptrain,batch_size) 65 | 66 | def read_labels(self, file_name): 67 | """ 68 | file_name:the byte file's direction 69 | """ 70 | label_file=open(file_name,'rb') 71 | print(label_file) 72 | # get the basic information about the labels 73 | label_file.seek(0) 74 | magic_number = label_file.read(4) 75 | magic_number = struct.unpack('>i', magic_number) 76 | print('Magic Number: ' + str(magic_number[0])) 77 | 78 | data_type = label_file.read(4) 79 | data_type = struct.unpack('>i', data_type) 80 | print('Number of Lables: ' + str(data_type[0])) 81 | 82 | labels = [] 83 | for idx in range(data_type[0]): 84 | label_file.seek(8 + idx) 85 | tmp_d = label_file.read(1) 86 | tmp_d = struct.unpack('>B', tmp_d) 87 | labels.append(tmp_d) 88 | return np.array(labels) 89 | 90 | 91 | def read_images(self, file_name): 92 | """ 93 | file_name:the byte file's direction 94 | """ 95 | img_file = open(file_name, 'rb') 96 | print(img_file) 97 | # get the basic information about the images 98 | img_file.seek(0) 99 | magic_number = img_file.read(4) 100 | magic_number = struct.unpack('>i', magic_number) 101 | print('Magic Number: ' + str(magic_number[0])) 102 | 103 | data_type = img_file.read(4) 104 | data_type = struct.unpack('>i', data_type) 105 | print('Number of Images: ' + str(data_type[0])) 106 | 107 | dim = img_file.read(8) 108 | dimr = struct.unpack('>i', dim[0:4]) 109 | dimr = dimr[0] 110 | print('Number of Rows: ' + str(dimr)) 111 | dimc = struct.unpack('>i', dim[4:]) 112 | dimc = dimc[0] 113 | print('Number of Columns:' + str(dimc)) 114 | 115 | 116 | images=[] 117 | for idx in range(data_type[0]): 118 | image = np.ndarray(shape=(dimr, dimc)) 119 | img_file.seek(16 + dimc * dimr * idx) 120 | 121 | for row in range(dimr): 122 | for col in range(dimc): 123 | tmp_d = img_file.read(1) 124 | tmp_d = struct.unpack('>B', tmp_d) 125 | image[row, col] = tmp_d[0] 126 | images.append(image) 127 | return np.array(images) 128 | -------------------------------------------------------------------------------- /Datasets/wider.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import h5py 3 | import numpy as np 4 | import cPickle as pickle 5 | import os 6 | import time 7 | 8 | class WiderDataset(): 9 | annotation_names=['blur','pose','occlusion', 10 | 'invalid','expression','illumination'] 11 | def __init__(self,path,set='val'): 12 | self.set=set 13 | self._data_path=path 14 | self.cache_file=os.path.join(self._data_path,'cache',self.set+'.pth') 15 | if not self._load_cache(): 16 | self.file_names,self.boxes,self.annotations=self._read_data() 17 | pickle.dump([self.file_names,self.boxes,self.annotations],open(self.cache_file,'w')) 18 | self.len=len(self.file_names) 19 | 20 | def _show_pic_boxes(self,im,boxes,delay=1000): 21 | for box in boxes: 22 | x0,y0,x1,y1=[int(i) for i in np.round(box)] 23 | cv2.rectangle(im,(x0,y0),(x1,y1),(255,0,0)) 24 | cv2.imshow("draw_pic", im) 25 | cv2.waitKey(delay) 26 | 27 | def show_pic_boxes(self,idx): 28 | im=cv2.imread(self.file_names[idx]) 29 | self._show_pic_boxes(im,self.boxes[idx]) 30 | 31 | 32 | def _load_cache(self): 33 | if os.path.exists(self.cache_file): 34 | print 'load from cache file',self.cache_file 35 | st_time=time.time() 36 | self.file_names, self.boxes, self.annotations =pickle.load(open(self.cache_file)) 37 | print time.time()-st_time,'s used' 38 | return True 39 | return False 40 | 41 | def _read_data(self): 42 | mat_path=self._data_path + '/wider_face_split/wider_face_%s.mat'%(self.set,) 43 | print 'load from mat file', mat_path 44 | st_time = time.time() 45 | f = h5py.File(mat_path) 46 | # f=sio.loadmat(self._data_path + '/wider_face_split/wider_face_val.mat') 47 | annotations = {'blur':[],'pose':[],'occlusion':[], 48 | 'invalid':[],'expression':[],'illumination':[]} 49 | boxes=[] 50 | file_names = [] 51 | for name in self.annotation_names: 52 | for folder in f[name+'_label_list'][0]: 53 | for image in f[folder][0]: 54 | for value in f[image]: 55 | annotations[name].append(value) 56 | for folder in f['face_bbx_list'][0]: 57 | for im_file in f[folder][0]: 58 | x = np.array([bbxs for bbxs in f[im_file]]) 59 | x = x.transpose() 60 | x[:, 2:] += x[:, :2] 61 | boxes.append(x) 62 | for folder in f['file_list'][0]: 63 | for im_file in f[folder][0]: 64 | s = "".join([chr(c) for c in f[im_file]]) 65 | file_names.append(self._data_path + '/images_no_fold/' + s + '.jpg') 66 | print time.time() - st_time, 's used' 67 | return file_names, boxes, annotations 68 | 69 | def sift_hard(self,idx,min_area=240,blur=True,pose=True,occlusion=True,invalid=True): 70 | # area = 0 for no area sift 71 | reserve=np.ones(len(self.boxes[idx])) 72 | if min_area: 73 | area=np.prod(self.boxes[idx][:,2:]-self.boxes[idx][:,:2],1) 74 | # print area 75 | reserve[area < min_area] = 0 76 | # if max_area: 77 | # area = np.prod(self.boxes[idx][:, 2:] - self.boxes[idx][:, :2], 1) 78 | # reserve[area > max_area] = 0 79 | if blur: 80 | reserve[self.annotations['blur'][idx]!=0]=0 81 | if pose: 82 | reserve[self.annotations['pose'][idx]!=0]=0 83 | if occlusion: 84 | reserve[self.annotations['occlusion'][idx]!=0]=0 85 | if invalid: 86 | reserve[self.annotations['invalid'][idx] != 0] = 0 87 | return reserve 88 | 89 | def get_idx_annotation(self,idx): 90 | annotation={} 91 | for name in self.annotation_names: 92 | annotation[name]=self.annotations[name][idx] 93 | return annotation 94 | 95 | def __len__(self): 96 | return self.len 97 | 98 | def __getitem__(self,idx): 99 | # given a idx then read the idxth image in wider face dataset 100 | # return [a numpy image with [1,height,width,BGR],boxes array ,annotations dict] 101 | im = cv2.imread(self.file_names[idx]) 102 | while im is None: 103 | idx-=1 104 | im = cv2.imread(self.file_names[idx]) 105 | boxes=self.boxes[idx] 106 | annotation=self.get_idx_annotation(idx) 107 | return im,boxes,annotation -------------------------------------------------------------------------------- /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()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/MxnetA.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | import mxnet as mx 3 | import mxnet.symbol as sym 4 | import json 5 | from analysis.layers import * 6 | import re 7 | import ctypes 8 | 9 | from mxnet.ndarray import NDArray 10 | import mxnet.ndarray as nd 11 | from mxnet.base import NDArrayHandle, py_str 12 | 13 | 14 | 15 | blob_dict=[] 16 | tracked_layers = [] 17 | 18 | def tmpnet(): 19 | x=sym.Variable('data') 20 | y=sym.Convolution(x,kernel=(3,3),num_filter=32) 21 | y=sym.Activation(y,'relu') 22 | y = sym.Convolution(y, kernel=(3, 3), num_filter=64,stride=(2,2),num_group=2) 23 | y=sym.softmax(y) 24 | return y 25 | 26 | def analyse(data_infos,module_json,data_name='data'): 27 | 28 | datas={} 29 | for info in data_infos: 30 | datas[info[1]]=info[2] 31 | nodes=json.loads(module_json)['nodes'] 32 | input=[] 33 | out=None 34 | for node in nodes: 35 | name=node['name'] 36 | bottoms=[str(nodes[i[0]]['name']) for i in node['inputs']] 37 | for i,bottom in enumerate(bottoms): 38 | if bottom+'_output' in datas: 39 | bottoms[i]=datas[bottom+'_output'] 40 | elif bottom+'_0' in datas: 41 | bottoms[i]=datas[bottom+'_0'] 42 | elif bottom in datas: 43 | bottoms[i]=datas[bottom] 44 | else: 45 | cur_node=node 46 | while True: 47 | bottom = [str(nodes[inp[0]]['name']) for inp in cur_node['inputs']][0] 48 | if bottom + '_output' in datas: 49 | bottoms[i] = datas[bottom + '_output'] 50 | break 51 | elif bottom + '_0' in datas: 52 | bottoms[i] = datas[bottom + '_0'] 53 | break 54 | elif bottom in datas: 55 | bottoms[i] = datas[bottom] 56 | break 57 | try: 58 | bottom_node = nodes[cur_node['inputs'][0][0]] 59 | except: 60 | pass 61 | cur_node=bottom_node 62 | if data_name==name: 63 | input.append(Blob(datas[data_name])) 64 | elif node['op']=='Convolution': 65 | kernel=eval(node['attrs']['kernel']) 66 | num_out=eval(node['attrs']['num_filter']) 67 | group_size=eval(node['attrs'].get('num_group','1')) 68 | pad=eval(node['attrs'].get('pad','(0,0)')) 69 | stride=eval(node['attrs'].get('stride','(1,1)')) 70 | x=Blob(bottoms[0]) 71 | out=Conv(x,kernel_size=kernel,stride=stride,pad=pad, 72 | num_out=num_out,group_size=group_size,name=name) 73 | tracked_layers.append(out) 74 | elif node['op']=='BatchNorm': 75 | x=Blob(bottoms[0]) 76 | out = Norm(x, 'batch_norm',name=name) 77 | tracked_layers.append(out) 78 | elif node['op']=='FullyConnected': 79 | x=Blob(bottoms[0]) 80 | num_hidden=eval(node['attrs']['num_hidden']) 81 | out=Fc(x,num_hidden,name=name) 82 | tracked_layers.append(out) 83 | elif node['op']=='Activation': 84 | pass 85 | elif 'elemwise' in node['op']: 86 | pass 87 | 88 | 89 | class Monitor(object): 90 | def __init__(self, interval=1, pattern='.*', sort=False): 91 | def stat(x): 92 | return x.shape 93 | self.stat_func = stat 94 | self.interval = interval 95 | self.activated = False 96 | self.queue = [] 97 | self.step = 0 98 | self.exes = [] 99 | self.re_prog = re.compile(pattern) 100 | self.sort = sort 101 | def stat_helper(name, array): 102 | array = ctypes.cast(array, NDArrayHandle) 103 | array = NDArray(array, writable=False) 104 | if not self.activated or not self.re_prog.match(py_str(name)): 105 | return 106 | self.queue.append((self.step, py_str(name), stat(array))) 107 | self.stat_helper = stat_helper 108 | 109 | def install(self, exe): 110 | exe.set_monitor_callback(self.stat_helper) 111 | self.exes.append(exe) 112 | 113 | def tic(self): 114 | if self.step % self.interval == 0: 115 | for exe in self.exes: 116 | for array in exe.arg_arrays: 117 | array.wait_to_read() 118 | for array in exe.aux_arrays: 119 | array.wait_to_read() 120 | self.queue = [] 121 | self.activated = True 122 | self.step += 1 123 | 124 | def toc(self): 125 | if not self.activated: 126 | return [] 127 | for exe in self.exes: 128 | for array in exe.arg_arrays: 129 | array.wait_to_read() 130 | for array in exe.aux_arrays: 131 | array.wait_to_read() 132 | for exe in self.exes: 133 | for name, array in zip(exe._symbol.list_arguments(), exe.arg_arrays): 134 | self.queue.append((self.step, name, self.stat_func(array))) 135 | for name, array in zip(exe._symbol.list_auxiliary_states(), exe.aux_arrays): 136 | # if self.re_prog.match(name): 137 | self.queue.append((self.step, name, self.stat_func(array))) 138 | self.activated = False 139 | res = [] 140 | if self.sort: 141 | self.queue.sort(key=lambda x: x[1]) 142 | for n, k, v_list in self.queue: 143 | res.append((n, k, v_list)) 144 | self.queue = [] 145 | return res 146 | def toc_print(self): 147 | pass 148 | 149 | def profiling_symbol(symbol,data_shape,data_name='data'): 150 | monitor = Monitor() 151 | model=mx.mod.Module(symbol) 152 | model.bind(data_shapes=[(data_name,tuple(data_shape))]) 153 | model.install_monitor(monitor) 154 | model.init_params() 155 | monitor.tic() 156 | model.forward(mx.io.DataBatch(data=(nd.ones(data_shape),))) 157 | data_infos=monitor.toc() 158 | module_json=symbol.tojson() 159 | analyse(data_infos,module_json,data_name) 160 | -------------------------------------------------------------------------------- /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 | 13 | def _analyse(module,raw_input): 14 | input=[] 15 | for i in raw_input: 16 | s = i.size() 17 | if len(s)==4: 18 | #input.append(Blob([s[0],s[2],s[3],s[1]])) 19 | input.append(Blob([s[0],s[1],s[2],s[3]])) 20 | else: 21 | input.append(Blob(s)) 22 | out=None 23 | if isinstance(module,nn.Conv2d): 24 | out=Conv(input[0],module.kernel_size,module.out_channels, 25 | module.stride,module.padding,group_size=module.groups) 26 | elif isinstance(module,nn.BatchNorm2d): 27 | out=Norm(input[0],'batch_norm') 28 | elif isinstance(module,nn.Linear): 29 | out=fc(input[0],module.out_features) 30 | elif isinstance(module,nn.MaxPool2d): 31 | out = pool(input[0], module.kernel_size,module.stride,module.padding, 32 | name='max_pool',pool_type='max') 33 | elif isinstance(module,nn.AvgPool2d): 34 | out = pool(input[0], module.kernel_size,module.stride,module.padding, 35 | name='avg_pool',pool_type='avg') 36 | elif isinstance(module,nn.ReLU): 37 | out = Activation(input[0],'relu') 38 | if out: 39 | tracked_layers.append(out) 40 | else: 41 | print('WARNING: skip Module {}' .format(module)) 42 | 43 | def module_hook(module, input, output): 44 | # print('module hook') 45 | # print module 46 | # for i in input: 47 | # print ('input',i.size()) 48 | # for i in output: 49 | # print('out', i.size()) 50 | _analyse(module,input) 51 | 52 | def register(module): 53 | module.register_forward_hook(module_hook) 54 | 55 | def analyse(net, inputs): 56 | """ 57 | analyse the network given input 58 | :param net: torch.nn.Module 59 | :param inputs: torch.Variable, torch.Tensor or list of them 60 | :return: blob_dict, tracked_layers 61 | """ 62 | del tracked_layers[:] 63 | del blob_dict[:] 64 | if inputs is not list: 65 | raw_inputs=[inputs] 66 | _inputs=[] 67 | for i in raw_inputs: 68 | if isinstance(i,Variable): 69 | _inputs.append(i) 70 | elif isinstance(i,torch.Tensor): 71 | _inputs.append(Variable(i)) 72 | elif isinstance(i,np.ndarray): 73 | _inputs.append(Variable(torch.Tensor(i))) 74 | else: 75 | raise NotImplementedError("Not Support the input type {}".format(type(i))) 76 | net.apply(register) 77 | net.forward(*_inputs) 78 | return blob_dict,tracked_layers 79 | 80 | def profilling(net,input): 81 | """ Old API of analyse """ 82 | 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/blob.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | # the blob with shape of (h,w,c) or (batch,h,w,c) for image 4 | class Blob(): 5 | def __init__(self,shape,father=None): 6 | shape=[int(i) for i in shape] 7 | self._data=None 8 | self.shape=[int(i) for i in list(shape)] 9 | self.father=type(father)==list and father or [father] 10 | 11 | @property 12 | def data(self): 13 | raise NotImplementedError('Blob.data is removed from this version of nn_tools, you should use .shape') 14 | 15 | @property 16 | def size(self): 17 | return np.prod(self.shape) 18 | 19 | @property 20 | def w(self): 21 | return self.shape[2] 22 | @property 23 | def h(self): 24 | return self.shape[1] 25 | 26 | @property 27 | def c(self): 28 | return self.shape[3] 29 | 30 | @property 31 | def batch_size(self): 32 | return self.shape[0] 33 | 34 | @property 35 | def dim(self): 36 | return len(self.shape) 37 | 38 | def new(self,father): 39 | return Blob(self.shape,father) 40 | def __getitem__(self, key): 41 | return self.shape[key] 42 | def __str__(self): 43 | return str(self.shape) 44 | def flaten(self): 45 | 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): 107 | # input is the instance of blob.Blob with shape (h,w,c) or (batch,h,w,c) 108 | super(Sliding,self).__init__(input,name=name) 109 | if self.input.dim!=4: 110 | raise ValueError('Sliding must have a input with (batch,w,h,c)') 111 | self.input_w = self.input.w 112 | self.input_h = self.input.h 113 | self.batch_size = self.input.batch_size 114 | self.in_channel = self.input.c 115 | 116 | if type(kernel_size)==int: 117 | self.kernel_size=[kernel_size,kernel_size] 118 | else: 119 | self.kernel_size=[i for i in kernel_size] 120 | if len(self.kernel_size)==1:self.kernel_size*=2 121 | if type(stride)==int: 122 | self.stride=[stride,stride] 123 | else: 124 | self.stride=[i for i in stride] 125 | if len(self.stride)==1:self.stride*=2 126 | elif len(self.stride)==0:self.stride=[1,1] 127 | elif len(self.stride)>2:raise AttributeError 128 | if type(pad)==int: 129 | self.pad=[pad,pad] 130 | else: 131 | self.pad=[i for i in pad] 132 | if len(self.pad)==1:self.pad*=2 133 | elif len(self.pad)==0:self.pad=[0,0] 134 | elif len(self.pad)>2:raise AttributeError 135 | self.num_out=num_out 136 | self.layer_info='kernel=%dx%d,stride=%dx%d,pad=%dx%d'%(self.kernel_size[0],self.kernel_size[1], 137 | self.stride[0],self.stride[1],self.pad[0],self.pad[1]) 138 | #calc out 139 | 140 | if not ceil: 141 | out_w=np.floor(float(self.input_w +self.pad[0]*2-self.kernel_size[0])/self.stride[0])+1 142 | out_h= np.floor(float(self.input_h + self.pad[1] * 2 - self.kernel_size[1]) / self.stride[1]) + 1 143 | else: 144 | out_w = np.ceil(float(self.input_w + self.pad[0] * 2 - self.kernel_size[0]) / self.stride[0]) + 1 145 | out_h = np.ceil(float(self.input_h + self.pad[1] * 2 - self.kernel_size[1]) / self.stride[1]) + 1 146 | self.out=Blob([self.batch_size,out_w,out_h,num_out],self) 147 | 148 | class Conv(Sliding): 149 | def __init__(self,input,kernel_size,num_out,stride=1,pad=0, 150 | activation='relu',name='conv',ceil=False,group_size=1): 151 | if isinstance(input,Base): 152 | input=input() 153 | Sliding.__init__(self,input,kernel_size,num_out,stride,pad,name=name,ceil=ceil) 154 | self.layer_info+=',num_out=%d'%(num_out) 155 | self.dot = np.prod(self.out.shape) * np.prod(self.kernel_size) * self.in_channel 156 | self.weight_size = np.prod(self.kernel_size) * num_out * self.in_channel 157 | if group_size!=1: 158 | self.layer_info += ',group_size=%d' % (group_size) 159 | self.dot /= group_size 160 | self.weight_size /= group_size 161 | self.add = self.dot 162 | if activation: 163 | Activation(self.out,activation) 164 | 165 | class Pool(Sliding): 166 | def __init__(self,input,kernel_size,stride=1,pad=0,name='pool',pool_type='max',ceil=False): 167 | # pool_type: 0 is max, 1 is avg/ave in Caffe 168 | if isinstance(input,Base): 169 | input=input() 170 | Sliding.__init__(self,input,kernel_size,input.c,stride,pad,name=name,ceil=ceil) 171 | self.pool_type=pool_type 172 | self.layer_info+=',type=%s'%(pool_type) 173 | if pool_type in ['max',0]: 174 | self.compare= np.prod(self.out.shape) * (np.prod(self.kernel_size) - 1) 175 | elif pool_type in ['avg','ave',1]: 176 | self.add = np.prod(self.input.shape) 177 | self.dot = np.prod(self.out.shape) 178 | else: 179 | print("WARNING, NOT IMPLEMENT POOL TYPE %s PROFILING at %s, CONTINUE"%(pool_type,name)) 180 | pool=Pool 181 | 182 | class InnerProduct(Base): 183 | def __init__(self,input,num_out,activation='relu',name='innerproduct'): 184 | if isinstance(input,Base): 185 | input=input() 186 | Base.__init__(self,input,name=name) 187 | self.left_dim=np.prod(input.shape[1:]) 188 | self.num_out=num_out 189 | self.dot=self.num_out*self.input_size 190 | self.add=self.num_out*self.input_size 191 | self.out=Blob([input[0],self.num_out],self) 192 | self.weight_size = self.num_out * self.left_dim 193 | if activation: 194 | Activation(self.out,activation) 195 | Fc=InnerProduct 196 | fc=InnerProduct 197 | 198 | class Permute(Base): 199 | def __init__(self, input,dims, name='permute'): 200 | super(Permute,self).__init__(input,name) 201 | self.out = Blob(dims,self) 202 | 203 | class Flatten(Base): 204 | def __init__(self,input, name='permute'): 205 | super(Flatten, self).__init__(input, name) 206 | dim=[np.prod(input.shape)] 207 | self.out = Blob(dim, self) 208 | 209 | class Eltwise(Base): 210 | def __init__(self,inputs,type='sum',name='eltwise'): 211 | super(Eltwise,self).__init__(inputs,name,) 212 | self.out=inputs[0].new(self) 213 | if type in ['sum','SUM']: 214 | self.add=np.prod(self.out.shape) 215 | elif type in ['product','PROD']: 216 | self.dot=np.prod(self.out.shape) 217 | elif type in ['max','MAX']: 218 | self.compare=np.prod(self.out.shape) 219 | else: 220 | raise AttributeError('the Eltwise layer type must be sum, max or product') 221 | 222 | class Slice(Base): 223 | def __init__(self,input,slice_point,axis,name='slice'): 224 | super(Slice,self).__init__(input,name,) 225 | self.out=[] 226 | last=0 227 | for p in slice_point: 228 | print(p,list(input.shape)) 229 | shape1=list(input.shape) 230 | shape1[axis] = p-last 231 | last=p 232 | self.out+=[Blob(shape1)] 233 | shape1 = list(input.shape) 234 | print(last,shape1,input.shape[axis]) 235 | shape1[axis] = input.shape[axis] - last 236 | self.out += [Blob(shape1)] 237 | 238 | class Reshape(Base): 239 | def __init__(self,input,shape,name='reshape'): 240 | super(Reshape,self).__init__(input,name) 241 | shape=list(shape) 242 | for i in range(len(shape)): 243 | if shape[i]==0: 244 | shape[i]=input.shape[i] 245 | self.out=Blob(shape) 246 | 247 | 248 | class Concat(Base): 249 | def __init__(self,inputs,axis,name='concat'): 250 | super(Concat,self).__init__(inputs,name,) 251 | outc=0 252 | for input in inputs: 253 | outc+=input[axis] 254 | self.out=Blob(inputs[0].shape,self) 255 | self.out.shape[axis]=outc 256 | 257 | class Scale(Base): 258 | def __init__(self, input, factor=None, name='scale'): 259 | super(Scale, self).__init__(input, name, ) 260 | self.out = input.new(self) 261 | 262 | self.dot=self.input_size 263 | # TODO scale analysis 264 | 265 | class Softmax(Base): 266 | def __init__(self, input, factor=None, name='softmax'): 267 | super(Softmax, self).__init__(input, name, ) 268 | self.out = input.new(self) 269 | self.power=self.input_size 270 | self.add=self.input_size 271 | self.dot=self.input_size 272 | self.layer_info="softmax" 273 | 274 | class Dropout(Base): 275 | def __init__(self,input,name='dropout'): 276 | if isinstance(input,Base): 277 | input=input() 278 | Base.__init__(self,input,name=name) 279 | 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, 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 | 37 | if csv_save_path!=None: 38 | with open(csv_save_path,'w') as file: 39 | writer=csv.writer(file) 40 | writer.writerow(save_items) 41 | for layer in print_list: 42 | writer.writerow(layer) 43 | if print_detail: 44 | sum[0] = 'SUM' 45 | print_list.append(sum) 46 | print_table(print_list,save_items) 47 | else: 48 | print_list=[] 49 | for idx,item in enumerate(sum): 50 | if item>0: 51 | if human_readable: 52 | print_list.append('%s:%s' % (save_items[idx], get_human_readable(item))) 53 | else: 54 | print_list.append('%s:%.3e'%(save_items[idx],item)) 55 | print(print_list) 56 | print('saved!') 57 | 58 | def get_layer_blox_from_blobs(blobs): 59 | layers=[] 60 | def creator_search(blob): 61 | for father in blob.father: 62 | if isinstance(father,Base) and father not in layers: 63 | layers.append(father) 64 | if father.muti_input==True: 65 | for input in father.input: 66 | creator_search(input) 67 | else: 68 | creator_search(father.input) 69 | for blob in blobs: 70 | creator_search(blob) 71 | return layers 72 | 73 | def print_table(datas,names): 74 | 75 | types=[] 76 | for i in datas[0]: 77 | try: 78 | i=int(float(i)) 79 | types.append('I') 80 | except: 81 | types.append('S') 82 | for l in datas: 83 | s='' 84 | for i,t in zip(l,types): 85 | if t=='I': 86 | 87 | i=int(float(i)) 88 | s+=('%.1E'%i).center(10) 89 | else: 90 | i=str(i) 91 | if len(i)>20: 92 | i=i[:17]+'...' 93 | s+=i.center(20) 94 | s+='|' 95 | print(s) 96 | s = '' 97 | for i,t in zip(names,types): 98 | 99 | if t == 'I': 100 | s += i.center(10) 101 | else: 102 | if len(i) > 20: 103 | i = i[:17] + '...' 104 | s += i.center(20) 105 | s += '|' 106 | print(s) 107 | 108 | def print_by_blob(blobs,print_items=('name', 'layer_info', 'input', 'out', 'dot', 'add', 'compare','ops', 'weight_size','activation_size')): 109 | layers=get_layer_blox_from_blobs(blobs) 110 | print_list = [] 111 | for layer in layers: 112 | print_list.append([str(getattr(layer, param)) for param in print_items]) 113 | pprint.pprint(print_list, depth=3, width=200) 114 | 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/MGN-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 option import args as configs 16 | from model import mgn 17 | from torchvision.models import resnet 18 | """ 19 | Supporting analyse the inheritors of torch.nn.Moudule class. 20 | 21 | Command：`pytorch_analyser.py [-h] [--out OUT] [--class_args ARGS] path class_name shape` 22 | - The path is the python file path which contaning your class. 23 | - The class_name is the class name in your python file. 24 | - 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. 25 | - The out (optinal) is path to save the csv file, default is '/tmp/pytorch_analyse.csv'. 26 | - The class_args (optional) is the args to init the class in python file, default is empty. 27 | 28 | For example `python pytorch_analyser.py tmp/pytorch_analysis_test.py ResNet218 1,3,224,224` 29 | """ 30 | 31 | 32 | if __name__=="__main__": 33 | parser=argparse.ArgumentParser() 34 | parser.add_argument('--path',help='python file location',default = "MGN_analysis_example.py",type=str) 35 | parser.add_argument('--name',help='the class name or instance name in your python file',default = "mgn",type=str) 36 | parser.add_argument('--shape',help='input shape of the network(split by comma `,`), image shape should be: batch,c,h,w',type=str) 37 | parser.add_argument('--out',help='path to save the csv file',default='pytorch_analyse.csv',type=str) 38 | parser.add_argument('--class_args',help='args to init the class in python file',default='net',type=str) 39 | 40 | args=parser.parse_args() 41 | # net = mgn.MGN(configs) 42 | # # print(net) 43 | # # net.load('/home/wdd/Work',) 44 | # # net = mgn.MGN() 45 | 46 | # # state_dict = torch.load('/home/wdd/Work/Pytorch/pytorch-caffe-darknet-convert-master/torch_model/model_160_max.pt') 47 | # checkpoint = torch.load('/home/shining/Projects/github-projects/pytorch-project/nn_tools/model_100.pt') 48 | # net.load_state_dict(checkpoint) 49 | # name = 'MGN' 50 | # # net = inception_v3(True, transform_input=False) 51 | # net.eval() 52 | name='resnet18' 53 | net=resnet.resnet18() 54 | net.eval() 55 | #input=Variable(torch.ones([1,3,224,224])) 56 | 57 | 58 | x = torch.rand( [1,3,224, 224]) 59 | blob_dict, layers = analyse(net, x) 60 | save_csv(layers, args.out) 61 | 62 | 63 | -------------------------------------------------------------------------------- /example/MGN_2_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 | from model import mgn 8 | from option import args 9 | import utils.utility as utility 10 | from torchvision.models import resnet 11 | 12 | 13 | if __name__=='__main__': 14 | # ckpt = torch.load(self.dir + '/map_log.pt') 15 | net = mgn.MGN(args) 16 | # print(net) 17 | # net.load('/home/wdd/Work',) 18 | # net = mgn.MGN() 19 | 20 | # state_dict = torch.load('/home/wdd/Work/Pytorch/pytorch-caffe-darknet-convert-master/torch_model/model_160_max.pt') 21 | checkpoint = torch.load('/home/shining/Projects/github-projects/pytorch-project/nn_tools/model_100.pt') 22 | net.load_state_dict(checkpoint) 23 | 24 | 25 | # net = Model(num_classes=2220) 26 | # # print ('person_ReID:', m) 27 | # state_dict = torch.load('/home/wdd/Work/Pytorch/pytorch-caffe-darknet-convert-master/torch_model/Ep600_ckpt.pth') 28 | # # print ("state_dict: ", state_dict) 29 | # state_dict = state_dict["state_dicts"][0] 30 | # net.load_state_dict(state_dict) 31 | 32 | name = 'MGN' 33 | # net = inception_v3(True, transform_input=False) 34 | net.eval() 35 | 36 | input = torch.ones([1, 3, 384, 128]) 37 | out = net.forward(input) 38 | pytorch_to_caffe.trans_net(net, input, name) 39 | pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name)) 40 | pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name)) -------------------------------------------------------------------------------- /example/MGN_analysis_example.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.path.insert(0,'.') 3 | import torch 4 | import torch.nn as nn 5 | from torchvision.models import resnet 6 | import pytorch_analyser 7 | from option import args 8 | from model import mgn 9 | 10 | if __name__=='__main__': 11 | # ckpt = torch.load(self.dir + '/map_log.pt') 12 | net = mgn.MGN(args) 13 | # print(net) 14 | # net.load('/home/wdd/Work',) 15 | # net = mgn.MGN() 16 | 17 | # state_dict = torch.load('/home/wdd/Work/Pytorch/pytorch-caffe-darknet-convert-master/torch_model/model_160_max.pt') 18 | checkpoint = torch.load('/home/shining/Projects/github-projects/pytorch-project/nn_tools/model_100.pt') 19 | net.load_state_dict(checkpoint) 20 | 21 | 22 | # net = Model(num_classes=2220) 23 | # # print ('person_ReID:', m) 24 | # state_dict = torch.load('/home/wdd/Work/Pytorch/pytorch-caffe-darknet-convert-master/torch_model/Ep600_ckpt.pth') 25 | # # print ("state_dict: ", state_dict) 26 | # state_dict = state_dict["state_dicts"][0] 27 | # net.load_state_dict(state_dict) 28 | 29 | name = 'MGN' 30 | # net = inception_v3(True, transform_input=False) 31 | net.eval() 32 | input_tensor=torch.ones(1,3,384,128) 33 | blob_dict, tracked_layers=pytorch_analyser.analyse(net,input_tensor) 34 | pytorch_analyser.save_csv(tracked_layers,'/tmp/analysis.csv') 35 | 36 | -------------------------------------------------------------------------------- /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=Variable(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/option.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | parser = argparse.ArgumentParser(description='MGN') 4 | 5 | parser.add_argument('--nThread', type=int, default=1, help='number of threads for data loading') 6 | parser.add_argument('--cpu', default='True', help='use cpu only') 7 | parser.add_argument('--nGPU', type=int, default=1, help='number of GPUs') 8 | 9 | # parser.add_argument("--datadir", type=str, default="dataset_4", help='dataset directory') 10 | # parser.add_argument("--datadir", type=str, default="market_cuhk03_duke", help='dataset directory') 11 | parser.add_argument("--datadir", type=str, default="/home/wdd/Work/Pytorch/MGN-pytorch-master/Market-1501-v15.09.15", help='dataset directory') 12 | parser.add_argument('--data_train', type=str, default='Market1501', help='train dataset name') 13 | parser.add_argument('--data_test', type=str, default='Market1501', help='test dataset name') 14 | 15 | parser.add_argument('--reset', action='store_true', help='reset the training') 16 | parser.add_argument("--epochs", type=int, default=80, help='number of epochs to train') 17 | parser.add_argument('--test_every', type=int, default=20, help='do test per every N epochs') 18 | parser.add_argument("--batchid", type=int, default=8, help='the batch for id') 19 | parser.add_argument("--batchimage", type=int, default=4, help='the batch of per id') 20 | parser.add_argument("--batchtest", type=int, default=8, help='input batch size for test') 21 | # parser.add_argument('--test_only', action='store_true', help='set this option to test the model') 22 | parser.add_argument('--test_only', default='True', help='set this option to test the model') 23 | 24 | parser.add_argument('--model', default='MGN', help='model name') 25 | parser.add_argument('--loss', type=str, default='1*CrossEntropy+1*Triplet', help='loss function configuration') 26 | 27 | parser.add_argument('--act', type=str, default='relu', help='activation function') 28 | parser.add_argument('--pool', type=str, default='max', help='pool function') 29 | # parser.add_argument('--pool', type=str, default='avg', help='pool function') #ori 30 | parser.add_argument('--feats', type=int, default=256, help='number of feature maps') 31 | parser.add_argument('--height', type=int, default=384, help='height of the input image') 32 | parser.add_argument('--width', type=int, default=128, help='width of the input image') 33 | parser.add_argument('--num_classes', type=int, default=751, help='') 34 | 35 | 36 | parser.add_argument("--lr", type=float, default=2e-5, help='learning rate') 37 | parser.add_argument('--optimizer', default='ADAM', choices=('SGD','ADAM','ADAMAX','RMSprop'), help='optimizer to use (SGD | ADAM | ADAMAX | RMSprop)') 38 | parser.add_argument('--momentum', type=float, default=0.9, help='SGD momentum') 39 | parser.add_argument('--dampening', type=float, default=0.0005, help='SGD dampening') 40 | parser.add_argument('--nesterov', action='store_true', help='SGD nesterov') 41 | parser.add_argument('--beta1', type=float, default=0.9, help='ADAM beta1') 42 | parser.add_argument('--beta2', type=float, default=0.999, help='ADAM beta2') 43 | parser.add_argument('--amsgrad', action='store_true', help='ADAM amsgrad') 44 | parser.add_argument('--epsilon', type=float, default=1e-8, help='ADAM epsilon for numerical stability') 45 | parser.add_argument('--gamma', type=float, default=0.1, help='learning rate decay factor for step decay') 46 | parser.add_argument('--weight_decay', type=float, default=2e-4, help='weight decay') 47 | parser.add_argument('--decay_type', type=str, default='step', help='learning rate decay type') 48 | parser.add_argument('--lr_decay', type=int, default=60, help='learning rate decay per N epochs') 49 | 50 | parser.add_argument("--margin", type=float, default=1.2, help='') 51 | # parser.add_argument("--re_rank", action='store_true', help='') 52 | parser.add_argument("--re_rank", default='False', help='') 53 | parser.add_argument("--random_erasing", action='store_true', help='') 54 | parser.add_argument("--probability", type=float, default=0.5, help='') 55 | 56 | parser.add_argument("--savedir", type=str, default='saved_models', help='directory name to save') 57 | parser.add_argument("--outdir", type=str, default='out', help='') # not used 58 | parser.add_argument("--resume", type=int, default=100, help='resume from specific checkpoint') 59 | ####changed =========== 60 | parser.add_argument('--save', type=str, default='test', help='file name to save') 61 | parser.add_argument('--load', type=str, default='', help='file name to load') 62 | parser.add_argument('--save_models', action='store_true', help='save all intermediate models') 63 | parser.add_argument('--pre_train', type=str, default='', help='pre-trained model directory') 64 | 65 | args = parser.parse_args() 66 | 67 | for arg in vars(args): 68 | if vars(args)[arg] == 'True': 69 | print(arg,'is True') 70 | vars(args)[arg] = True 71 | elif vars(args)[arg] == 'False': 72 | vars(args)[arg] = False 73 | print(arg,'is False') 74 | 75 | -------------------------------------------------------------------------------- /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 | resnet18.eval() 12 | input=Variable(torch.ones([1,3,224,224])) 13 | #input=torch.ones([1,3,224,224]) 14 | pytorch_to_caffe.trans_net(resnet18,input,name) 15 | pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name)) 16 | 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 | # 2017.12.16 by xiaohang 2 | import sys 3 | from caffenet import * 4 | import numpy as np 5 | import argparse 6 | import torch.nn as nn 7 | from torch.autograd import Variable 8 | from torch.nn.parameter import Parameter 9 | import time 10 | 11 | #caffe load formate 12 | def load_image_caffe(imgfile): 13 | image = caffe.io.load_image(imgfile) 14 | transformer = caffe.io.Transformer({'data': (1, 3, args.height, args.width)}) 15 | transformer.set_transpose('data', (2, 0, 1)) 16 | transformer.set_mean('data', np.array([args.meanB, args.meanG, args.meanR])) 17 | transformer.set_raw_scale('data', args.scale) 18 | transformer.set_channel_swap('data', (2, 1, 0)) 19 | 20 | image = transformer.preprocess('data', image) 21 | image = image.reshape(1, 3, args.height, args.width) 22 | return image 23 | 24 | 25 | def load_synset_words(synset_file): 26 | lines = open(synset_file).readlines() 27 | synset_dict = dict() 28 | for i, line in enumerate(lines): 29 | synset_dict[i] = line.strip() 30 | return synset_dict 31 | 32 | def forward_pytorch(protofile, weightfile, image): 33 | net = CaffeNet(protofile, width=args.width, height=args.height, omit_data_layer=True, phase='TEST') 34 | if args.cuda: 35 | net.cuda() 36 | print(net) 37 | net.load_weights(weightfile) 38 | net.eval() 39 | image = torch.from_numpy(image) 40 | if args.cuda: 41 | image = Variable(image.cuda()) 42 | else: 43 | image = Variable(image) 44 | t0 = time.time() 45 | blobs = net(image) 46 | t1 = time.time() 47 | return t1-t0, blobs, net.models 48 | 49 | # Reference from: 50 | def forward_caffe(protofile, weightfile, image): 51 | if args.cuda: 52 | caffe.set_device(0) 53 | caffe.set_mode_gpu() 54 | else: 55 | caffe.set_mode_cpu() 56 | net = caffe.Net(protofile, weightfile, caffe.TEST) 57 | net.blobs['data'].reshape(1, 3, args.height, args.width) 58 | net.blobs['data'].data[...] = image 59 | t0 = time.time() 60 | output = net.forward() 61 | t1 = time.time() 62 | return t1-t0, net.blobs, net.params 63 | 64 | if __name__ == '__main__': 65 | parser = argparse.ArgumentParser(description='convert caffe to pytorch') 66 | parser.add_argument('--protofile', default='', type=str) 67 | parser.add_argument('--weightfile', default='', type=str) 68 | parser.add_argument('--imgfile', default='', type=str) 69 | parser.add_argument('--height', default=224, type=int) 70 | parser.add_argument('--width', default=224, type=int) 71 | parser.add_argument('--meanB', default=104, type=float) 72 | parser.add_argument('--meanG', default=117, type=float) 73 | parser.add_argument('--meanR', default=123, type=float) 74 | parser.add_argument('--scale', default=255, type=float) 75 | parser.add_argument('--synset_words', default='', type=str) 76 | parser.add_argument('--cuda', action='store_true', help='enables cuda') 77 | 78 | args = parser.parse_args() 79 | print(args) 80 | 81 | 82 | protofile = args.protofile 83 | weightfile = args.weightfile 84 | imgfile = args.imgfile 85 | 86 | image = load_image(imgfile) 87 | time_pytorch, pytorch_blobs, pytorch_models = forward_pytorch(protofile, weightfile, image) 88 | time_caffe, caffe_blobs, caffe_params = forward_caffe(protofile, weightfile, image) 89 | 90 | print('pytorch forward time %d', time_pytorch) 91 | print('caffe forward time %d', time_caffe) 92 | 93 | layer_names = pytorch_models.keys() 94 | blob_names = pytorch_blobs.keys() 95 | print('------------ Parameter Difference ------------') 96 | for layer_name in layer_names: 97 | if type(pytorch_models[layer_name]) in [nn.Conv2d, nn.Linear, Scale, Normalize]: 98 | pytorch_weight = pytorch_models[layer_name].weight.data 99 | if args.cuda: 100 | pytorch_weight = pytorch_weight.cpu().numpy() 101 | else: 102 | pytorch_weight = pytorch_weight.numpy() 103 | caffe_weight = caffe_params[layer_name][0].data 104 | weight_diff = abs(pytorch_weight - caffe_weight).sum() 105 | if type(pytorch_models[layer_name].bias) == Parameter: 106 | pytorch_bias = pytorch_models[layer_name].bias.data 107 | if args.cuda: 108 | pytorch_bias = pytorch_bias.cpu().numpy() 109 | else: 110 | pytorch_bias = pytorch_bias.numpy() 111 | caffe_bias = caffe_params[layer_name][1].data 112 | bias_diff = abs(pytorch_bias - caffe_bias).sum() 113 | print('%-30s weight_diff: %f bias_diff: %f' % (layer_name, weight_diff, bias_diff)) 114 | else: 115 | print('%-30s weight_diff: %f' % (layer_name, weight_diff)) 116 | elif type(pytorch_models[layer_name]) == nn.BatchNorm2d: 117 | if args.cuda: 118 | pytorch_running_mean = pytorch_models[layer_name].running_mean.cpu().numpy() 119 | pytorch_running_var = pytorch_models[layer_name].running_var.cpu().numpy() 120 | else: 121 | pytorch_running_mean = pytorch_models[layer_name].running_mean.numpy() 122 | pytorch_running_var = pytorch_models[layer_name].running_var.numpy() 123 | caffe_running_mean = caffe_params[layer_name][0].data/caffe_params[layer_name][2].data[0] 124 | caffe_running_var = caffe_params[layer_name][1].data/caffe_params[layer_name][2].data[0] 125 | running_mean_diff = abs(pytorch_running_mean - caffe_running_mean).sum() 126 | running_var_diff = abs(pytorch_running_var - caffe_running_var).sum() 127 | print('%-30s running_mean_diff: %f running_var_diff: %f' % (layer_name, running_mean_diff, running_var_diff)) 128 | 129 | print('------------ Output Difference ------------') 130 | for blob_name in blob_names: 131 | if args.cuda: 132 | pytorch_data = pytorch_blobs[blob_name].data.cpu().numpy() 133 | else: 134 | pytorch_data = pytorch_blobs[blob_name].data.numpy() 135 | caffe_data = caffe_blobs[blob_name].data 136 | diff = abs(pytorch_data - caffe_data).sum() 137 | print('%-30s pytorch_shape: %-20s caffe_shape: %-20s output_diff: %f' % (blob_name, pytorch_data.shape, caffe_data.shape, diff/pytorch_data.size)) 138 | 139 | if args.synset_words != '': 140 | print('------------ Classification ------------') 141 | synset_dict = load_synset_words(args.synset_words) 142 | if 'prob' in blob_names: 143 | if args.cuda: 144 | pytorch_prob = pytorch_blobs['prob'].data.cpu().view(-1).numpy() 145 | else: 146 | pytorch_prob = pytorch_blobs['prob'].data.view(-1).numpy() 147 | caffe_prob = caffe_blobs['prob'].data[0] 148 | print('pytorch classification top1: %f %s' % (pytorch_prob.max(), synset_dict[pytorch_prob.argmax()])) 149 | print('caffe classification top1: %f %s' % (caffe_prob.max(), synset_dict[caffe_prob.argmax()])) 150 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /keras_to_caffe.py: -------------------------------------------------------------------------------- 1 | import keras 2 | import Caffe.net as caffe 3 | import numpy as np 4 | from funcs import * 5 | 6 | def convert_filter(numpy_filter_weight): 7 | return np.transpose(numpy_filter_weight,(3,2,1,0)) 8 | 9 | def convert_fc(numpy_fc_weight): 10 | return np.transpose(numpy_fc_weight,(1,0)) 11 | 12 | def keras_weights_to_caffemodel(keras_model): 13 | """ 14 | Only Implement the conv layer and fc layer 15 | :param keras_model: 16 | :return: 17 | """ 18 | net=caffe.Net() 19 | layers=keras_model.layers 20 | 21 | for layer in layers: 22 | if type(layer)==keras.layers.Convolution2D: 23 | w,b=layer.get_weights() 24 | w=convert_filter(w) 25 | param=caffe.Layer_param(layer.name,'Convolution') 26 | net.add_layer_with_data(param,[w,b]) 27 | if type(layer)==keras.layers.Dense: 28 | w, b = layer.get_weights() 29 | w = convert_fc(w) 30 | param = caffe.Layer_param(layer.name, 'InnerProduct') 31 | net.add_layer_with_data(param, [w, b]) 32 | return net 33 | 34 | if __name__=='__main__': 35 | pass -------------------------------------------------------------------------------- /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/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 | -------------------------------------------------------------------------------- /mxnet_analyser.py: -------------------------------------------------------------------------------- 1 | # coding=utf-8 2 | from __future__ import absolute_import 3 | import argparse 4 | from analysis.MxnetA import * 5 | from analysis.utils import save_csv 6 | import os 7 | import sys 8 | import mxnet 9 | 10 | if __name__=="__main__": 11 | parser = argparse.ArgumentParser() 12 | parser.add_argument('path', help='python file location', type=str) 13 | parser.add_argument('name', help='the symbol object name or function that generate the symbol in your python file', type=str) 14 | parser.add_argument('shape', 15 | help='input shape of the network(split by comma `,`), image shape should be: batch,c,h,w', 16 | type=str) 17 | parser.add_argument('--out', help='path to save the csv file', default='/tmp/mxnet_analyse.csv', type=str) 18 | parser.add_argument('--func_args', help='args tuple parse to the function, eg. --func_args (3,"ABC")', default='', type=str) 19 | parser.add_argument('--func_kwargs', help='kwargs dict parse to the function, eg. --func_kwargs {a=1,c="OP"}', default='', type=str) 20 | 21 | args = parser.parse_args() 22 | path, filename = os.path.split(args.path) 23 | path=os.path.abspath(path) 24 | print(path) 25 | filename = os.path.splitext(filename)[0] 26 | sys.path.insert(0, path) 27 | exec ('from %s import %s as sym' % (filename, args.name)) 28 | if isinstance(sym, mxnet.sym.Symbol): 29 | sym = sym 30 | elif hasattr(sym,'__call__'): 31 | if args.func_kwargs!='': 32 | kwargs=eval(args.func_kwargs) 33 | else: 34 | kwargs={} 35 | if args.func_args!='': 36 | func_args=eval(args.func_args) 37 | else: 38 | func_args=[] 39 | sym = sym(*func_args,**kwargs) 40 | else: 41 | assert ("Error, The sym is not a instance of mxnet.sym.Symbol or function") 42 | shape = [int(i) for i in args.shape.split(',')] 43 | profiling_symbol(sym,shape) 44 | save_csv(tracked_layers, '/tmp/mxnet_analyse.csv') -------------------------------------------------------------------------------- /pytorch_analyse.csv: -------------------------------------------------------------------------------- 1 | name,layer_info,input,out,dot,add,compare,ops,weight_size,activation_size 2 | conv,"kernel=7x7,stride=2x2,pad=3x3,num_out=64","[1, 3, 224, 224]","[1, 112, 2, 64]",157351936,157351936,0,314703872,702464,14336 3 | batch_norm,None,"[1, 64, 112, 112]","[1, 64, 112, 112]",802816,802816,0,1605632,0,802816 4 | relu,None,"[1, 64, 112, 112]","[1, 64, 112, 112]",0,0,802816,802816,0,802816 5 | max_pool,"kernel=3x3,stride=2x2,pad=1x1,type=max","[1, 64, 112, 112]","[1, 56, 32, 112]",0,0,1605632,1605632,0,200704 6 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=64","[1, 64, 56, 56]","[1, 56, 64, 64]",115605504,115605504,0,231211008,32256,229376 7 | batch_norm,None,"[1, 64, 56, 56]","[1, 64, 56, 56]",200704,200704,0,401408,0,200704 8 | relu,None,"[1, 64, 56, 56]","[1, 64, 56, 56]",0,0,200704,200704,0,200704 9 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=64","[1, 64, 56, 56]","[1, 56, 64, 64]",115605504,115605504,0,231211008,32256,229376 10 | batch_norm,None,"[1, 64, 56, 56]","[1, 64, 56, 56]",200704,200704,0,401408,0,200704 11 | relu,None,"[1, 64, 56, 56]","[1, 64, 56, 56]",0,0,200704,200704,0,200704 12 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=64","[1, 64, 56, 56]","[1, 56, 64, 64]",115605504,115605504,0,231211008,32256,229376 13 | batch_norm,None,"[1, 64, 56, 56]","[1, 64, 56, 56]",200704,200704,0,401408,0,200704 14 | relu,None,"[1, 64, 56, 56]","[1, 64, 56, 56]",0,0,200704,200704,0,200704 15 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=64","[1, 64, 56, 56]","[1, 56, 64, 64]",115605504,115605504,0,231211008,32256,229376 16 | batch_norm,None,"[1, 64, 56, 56]","[1, 64, 56, 56]",200704,200704,0,401408,0,200704 17 | relu,None,"[1, 64, 56, 56]","[1, 64, 56, 56]",0,0,200704,200704,0,200704 18 | conv,"kernel=3x3,stride=2x2,pad=1x1,num_out=128","[1, 64, 56, 56]","[1, 28, 32, 128]",57802752,57802752,0,115605504,64512,114688 19 | batch_norm,None,"[1, 128, 28, 28]","[1, 128, 28, 28]",100352,100352,0,200704,0,100352 20 | relu,None,"[1, 128, 28, 28]","[1, 128, 28, 28]",0,0,100352,100352,0,100352 21 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=128","[1, 128, 28, 28]","[1, 28, 128, 128]",115605504,115605504,0,231211008,32256,458752 22 | batch_norm,None,"[1, 128, 28, 28]","[1, 128, 28, 28]",100352,100352,0,200704,0,100352 23 | conv,"kernel=1x1,stride=2x2,pad=0x0,num_out=128","[1, 64, 56, 56]","[1, 28, 32, 128]",6422528,6422528,0,12845056,7168,114688 24 | batch_norm,None,"[1, 128, 28, 28]","[1, 128, 28, 28]",100352,100352,0,200704,0,100352 25 | relu,None,"[1, 128, 28, 28]","[1, 128, 28, 28]",0,0,100352,100352,0,100352 26 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=128","[1, 128, 28, 28]","[1, 28, 128, 128]",115605504,115605504,0,231211008,32256,458752 27 | batch_norm,None,"[1, 128, 28, 28]","[1, 128, 28, 28]",100352,100352,0,200704,0,100352 28 | relu,None,"[1, 128, 28, 28]","[1, 128, 28, 28]",0,0,100352,100352,0,100352 29 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=128","[1, 128, 28, 28]","[1, 28, 128, 128]",115605504,115605504,0,231211008,32256,458752 30 | batch_norm,None,"[1, 128, 28, 28]","[1, 128, 28, 28]",100352,100352,0,200704,0,100352 31 | relu,None,"[1, 128, 28, 28]","[1, 128, 28, 28]",0,0,100352,100352,0,100352 32 | conv,"kernel=3x3,stride=2x2,pad=1x1,num_out=256","[1, 128, 28, 28]","[1, 14, 64, 256]",57802752,57802752,0,115605504,64512,229376 33 | batch_norm,None,"[1, 256, 14, 14]","[1, 256, 14, 14]",50176,50176,0,100352,0,50176 34 | relu,None,"[1, 256, 14, 14]","[1, 256, 14, 14]",0,0,50176,50176,0,50176 35 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=256","[1, 256, 14, 14]","[1, 14, 256, 256]",115605504,115605504,0,231211008,32256,917504 36 | batch_norm,None,"[1, 256, 14, 14]","[1, 256, 14, 14]",50176,50176,0,100352,0,50176 37 | conv,"kernel=1x1,stride=2x2,pad=0x0,num_out=256","[1, 128, 28, 28]","[1, 14, 64, 256]",6422528,6422528,0,12845056,7168,229376 38 | batch_norm,None,"[1, 256, 14, 14]","[1, 256, 14, 14]",50176,50176,0,100352,0,50176 39 | relu,None,"[1, 256, 14, 14]","[1, 256, 14, 14]",0,0,50176,50176,0,50176 40 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=256","[1, 256, 14, 14]","[1, 14, 256, 256]",115605504,115605504,0,231211008,32256,917504 41 | batch_norm,None,"[1, 256, 14, 14]","[1, 256, 14, 14]",50176,50176,0,100352,0,50176 42 | relu,None,"[1, 256, 14, 14]","[1, 256, 14, 14]",0,0,50176,50176,0,50176 43 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=256","[1, 256, 14, 14]","[1, 14, 256, 256]",115605504,115605504,0,231211008,32256,917504 44 | batch_norm,None,"[1, 256, 14, 14]","[1, 256, 14, 14]",50176,50176,0,100352,0,50176 45 | relu,None,"[1, 256, 14, 14]","[1, 256, 14, 14]",0,0,50176,50176,0,50176 46 | conv,"kernel=3x3,stride=2x2,pad=1x1,num_out=512","[1, 256, 14, 14]","[1, 7, 128, 512]",57802752,57802752,0,115605504,64512,458752 47 | batch_norm,None,"[1, 512, 7, 7]","[1, 512, 7, 7]",25088,25088,0,50176,0,25088 48 | relu,None,"[1, 512, 7, 7]","[1, 512, 7, 7]",0,0,25088,25088,0,25088 49 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=512","[1, 512, 7, 7]","[1, 7, 512, 512]",115605504,115605504,0,231211008,32256,1835008 50 | batch_norm,None,"[1, 512, 7, 7]","[1, 512, 7, 7]",25088,25088,0,50176,0,25088 51 | conv,"kernel=1x1,stride=2x2,pad=0x0,num_out=512","[1, 256, 14, 14]","[1, 7, 128, 512]",6422528,6422528,0,12845056,7168,458752 52 | batch_norm,None,"[1, 512, 7, 7]","[1, 512, 7, 7]",25088,25088,0,50176,0,25088 53 | relu,None,"[1, 512, 7, 7]","[1, 512, 7, 7]",0,0,25088,25088,0,25088 54 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=512","[1, 512, 7, 7]","[1, 7, 512, 512]",115605504,115605504,0,231211008,32256,1835008 55 | batch_norm,None,"[1, 512, 7, 7]","[1, 512, 7, 7]",25088,25088,0,50176,0,25088 56 | relu,None,"[1, 512, 7, 7]","[1, 512, 7, 7]",0,0,25088,25088,0,25088 57 | conv,"kernel=3x3,stride=1x1,pad=1x1,num_out=512","[1, 512, 7, 7]","[1, 7, 512, 512]",115605504,115605504,0,231211008,32256,1835008 58 | batch_norm,None,"[1, 512, 7, 7]","[1, 512, 7, 7]",25088,25088,0,50176,0,25088 59 | relu,None,"[1, 512, 7, 7]","[1, 512, 7, 7]",0,0,25088,25088,0,25088 60 | avg_pool,"kernel=7x7,stride=1x1,pad=0x0,type=avg","[1, 512, 7, 7]","[1, 1, 506, 7]",3542,25088,0,28630,0,3542 61 | innerproduct,None,"[1, 512]","[1, 1000]",512000,512000,0,1024000,512000,1000 62 | -------------------------------------------------------------------------------- /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 | from Caffe import caffe_net 4 | import torch.nn.functional as F 5 | from torch.autograd import Variable 6 | from Caffe import layer_param 7 | from torch.nn.modules.utils import _pair 8 | import numpy as np 9 | 10 | """ 11 | How to support a new layer type: 12 | layer_name=log.add_layer(layer_type_name) 13 | top_blobs=log.add_blobs() 14 | layer=caffe_net.Layer_param(xxx) 15 | 16 | [] 17 | log.cnet.add_layer(layer) 18 | """ 19 | 20 | # TODO: support the inplace output of the layers 21 | 22 | 23 | NET_INITTED=False 24 | class TransLog(object): 25 | def __init__(self): 26 | """ 27 | doing init() with inputs Variable before using it 28 | """ 29 | self.layers={} 30 | self.detail_layers={} 31 | self.detail_blobs={} 32 | self._blobs={} 33 | self._blobs_data=[] 34 | self.cnet=caffe_net.Caffemodel('') 35 | self.debug=True 36 | 37 | def init(self,inputs): 38 | """ 39 | :param inputs: is a list of input variables 40 | """ 41 | self.add_blobs(inputs) 42 | def add_layer(self,name='layer'): 43 | if name in self.layers: 44 | return self.layers[name] 45 | if name not in self.detail_layers.keys(): 46 | self.detail_layers[name] =0 47 | self.detail_layers[name] +=1 48 | name='{}{}'.format(name,self.detail_layers[name]) 49 | self.layers[name]=name 50 | if self.debug: 51 | print("{} was added to layers".format(self.layers[name])) 52 | return self.layers[name] 53 | 54 | def add_blobs(self, blobs,name='blob',with_num=True): 55 | rst=[] 56 | for blob in blobs: 57 | self._blobs_data.append(blob) # to block the memory address be rewrited 58 | blob=int(id(blob)) 59 | if name not in self.detail_blobs.keys(): 60 | self.detail_blobs[name] =0 61 | self.detail_blobs[name] +=1 62 | if with_num: 63 | rst.append('{}{}'.format(name,self.detail_blobs[name])) 64 | else: 65 | rst.append('{}'.format(name)) 66 | if self.debug: 67 | print("{}:{} was added to blobs".format(blob,rst[-1])) 68 | self._blobs[blob]=rst[-1] 69 | return rst 70 | def blobs(self, var): 71 | var=id(var) 72 | if self.debug: 73 | print("{}:{} getting".format(var, self._blobs[var])) 74 | try: 75 | return self._blobs[var] 76 | except: 77 | print("WARNING: CANNOT FOUND blob {}".format(var)) 78 | return None 79 | 80 | log=TransLog() 81 | 82 | def _conv2d(raw,input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): 83 | x=raw(input,weight,bias,stride,padding,dilation,groups) 84 | name=log.add_layer(name='conv') 85 | log.add_blobs([x],name='conv_blob') 86 | layer=caffe_net.Layer_param(name=name, type='Convolution', 87 | bottom=[log.blobs(input)], top=[log.blobs(x)]) 88 | layer.conv_param(x.size()[1],weight.size()[2:],stride=_pair(stride), 89 | pad=_pair(padding),dilation=_pair(dilation),bias_term=bias is not None) 90 | if bias is not None: 91 | layer.add_data(weight.cpu().data.numpy(),bias.cpu().data.numpy()) 92 | else: 93 | layer.param.convolution_param.bias_term=False 94 | layer.add_data(weight.cpu().data.numpy()) 95 | log.cnet.add_layer(layer) 96 | return x 97 | 98 | def _linear(raw,input, weight, bias=None): 99 | x=raw(input,weight,bias) 100 | layer_name=log.add_layer(name='fc') 101 | top_blobs=log.add_blobs([x],name='fc_blob') 102 | layer=caffe_net.Layer_param(name=layer_name,type='InnerProduct', 103 | bottom=[log.blobs(input)],top=top_blobs) 104 | layer.fc_param(x.size()[1]) 105 | if bias is not None: 106 | layer.add_data(weight.cpu().data.numpy(),bias.cpu().data.numpy()) 107 | else: 108 | layer.add_data(weight.cpu().data.numpy()) 109 | log.cnet.add_layer(layer) 110 | return x 111 | 112 | def _split(raw,tensor, split_size, dim=0): 113 | # split in pytorch is slice in caffe 114 | x=raw(tensor, split_size, dim) 115 | layer_name=log.add_layer('split') 116 | top_blobs=log.add_blobs(x,name='split_blob') 117 | layer=caffe_net.Layer_param(name=layer_name, type='Slice', 118 | bottom=[log.blobs(tensor)], top=top_blobs) 119 | slice_num=int(np.floor(tensor.size()[dim]/split_size)) 120 | slice_param=caffe_net.pb.SliceParameter(axis=dim,slice_point=[split_size*i for i in range(1,slice_num)]) 121 | layer.param.slice_param.CopyFrom(slice_param) 122 | log.cnet.add_layer(layer) 123 | return x 124 | 125 | 126 | def _pool(type,raw,input,x,kernel_size,stride,padding,ceil_mode): 127 | # TODO dilation,ceil_mode,return indices 128 | layer_name = log.add_layer(name='{}_pool'.format(type)) 129 | top_blobs = log.add_blobs([x], name='{}_pool_blob'.format(type)) 130 | layer = caffe_net.Layer_param(name=layer_name, type='Pooling', 131 | bottom=[log.blobs(input)], top=top_blobs) 132 | # TODO w,h different kernel, stride and padding 133 | # processing ceil mode 134 | layer.pool_param(kernel_size=kernel_size, stride=kernel_size if stride is None else stride, 135 | pad=padding, type=type.upper()) 136 | log.cnet.add_layer(layer) 137 | if ceil_mode==False and stride is not None: 138 | oheight = (input.size()[2] - _pair(kernel_size)[0] + 2 * _pair(padding)[0]) % (_pair(stride)[0]) 139 | owidth = (input.size()[3] - _pair(kernel_size)[1] + 2 * _pair(padding)[1]) % (_pair(stride)[1]) 140 | if oheight!=0 or owidth!=0: 141 | caffe_out=raw(input, kernel_size, stride, padding, ceil_mode=True) 142 | print("WARNING: the output shape miss match at {}: " 143 | 144 | "input {} output---Pytorch:{}---Caffe:{}\n" 145 | "This is caused by the different implementation that ceil mode in caffe and the floor mode in pytorch.\n" 146 | "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())) 147 | 148 | def _max_pool2d(raw,input, kernel_size, stride=None, padding=0, dilation=1, 149 | ceil_mode=False, return_indices=False): 150 | x = raw(input, kernel_size, stride, padding, dilation,ceil_mode, return_indices) 151 | _pool('max',raw,input, x, kernel_size, stride, padding,ceil_mode) 152 | return x 153 | 154 | def _avg_pool2d(raw,input, kernel_size, stride = None, padding = 0, ceil_mode = False, count_include_pad = True): 155 | x = raw(input, kernel_size, stride, padding, ceil_mode, count_include_pad) 156 | _pool('ave',raw,input, x, kernel_size, stride, padding,ceil_mode) 157 | return x 158 | 159 | def _max(raw,*args): 160 | x=raw(*args) 161 | if len(args)==1: 162 | # TODO max in one tensor 163 | assert NotImplementedError 164 | else: 165 | bottom_blobs=[] 166 | for arg in args: 167 | bottom_blobs.append(log.blobs(arg)) 168 | layer_name=log.add_layer(name='max') 169 | top_blobs=log.add_blobs([x],name='max_blob') 170 | layer=caffe_net.Layer_param(name=layer_name,type='Eltwise', 171 | bottom=bottom_blobs,top=top_blobs) 172 | layer.param.eltwise_param.operation =2 173 | log.cnet.add_layer(layer) 174 | return x 175 | 176 | def _cat(raw,inputs, dim=0): 177 | x=raw(inputs, dim) 178 | bottom_blobs=[] 179 | for input in inputs: 180 | bottom_blobs.append(log.blobs(input)) 181 | layer_name=log.add_layer(name='cat') 182 | top_blobs=log.add_blobs([x],name='cat_blob') 183 | layer=caffe_net.Layer_param(name=layer_name,type='Concat', 184 | bottom=bottom_blobs,top=top_blobs) 185 | layer.param.concat_param.axis =dim 186 | log.cnet.add_layer(layer) 187 | return x 188 | 189 | def _dropout(raw,input,p=0.5, training=False, inplace=False): 190 | x=raw(input,p, training, inplace) 191 | bottom_blobs=[log.blobs(input)] 192 | layer_name=log.add_layer(name='dropout') 193 | top_blobs=log.add_blobs([x],name=bottom_blobs[0],with_num=False) 194 | layer=caffe_net.Layer_param(name=layer_name,type='Dropout', 195 | bottom=bottom_blobs,top=top_blobs) 196 | layer.param.dropout_param.dropout_ratio = p 197 | layer.param.include.extend([caffe_net.pb.NetStateRule(phase=0)]) # 1 for test, 0 for train 198 | log.cnet.add_layer(layer) 199 | return x 200 | 201 | def _threshold(raw,input, threshold, value, inplace=False): 202 | # for threshold or relu 203 | if threshold==0 and value==0: 204 | x = raw(input,threshold, value, inplace) 205 | bottom_blobs=[log.blobs(input)] 206 | name = log.add_layer(name='relu') 207 | log.add_blobs([x], name='relu_blob') 208 | layer = caffe_net.Layer_param(name=name, type='ReLU', 209 | bottom=bottom_blobs, top=[log.blobs(x)]) 210 | log.cnet.add_layer(layer) 211 | return x 212 | if value!=0: 213 | raise NotImplemented("value !=0 not implemented in caffe") 214 | x=raw(input,input, threshold, value, inplace) 215 | bottom_blobs=[log.blobs(input)] 216 | layer_name=log.add_layer(name='threshold') 217 | top_blobs=log.add_blobs([x],name='threshold_blob') 218 | layer=caffe_net.Layer_param(name=layer_name,type='Threshold', 219 | bottom=bottom_blobs,top=top_blobs) 220 | layer.param.threshold_param.threshold = threshold 221 | log.cnet.add_layer(layer) 222 | return x 223 | 224 | def _prelu(raw, input, weight): 225 | # for threshold or prelu 226 | x = raw(input, weight) 227 | bottom_blobs=[log.blobs(input)] 228 | name = log.add_layer(name='prelu') 229 | log.add_blobs([x], name='prelu_blob') 230 | layer = caffe_net.Layer_param(name=name, type='PReLU', 231 | bottom=bottom_blobs, top=[log.blobs(x)]) 232 | if weight.size()[0]==1: 233 | layer.param.prelu_param.channel_shared=True 234 | layer.add_data(weight.cpu().data.numpy()[0]) 235 | else: 236 | layer.add_data(weight.cpu().data.numpy()) 237 | log.cnet.add_layer(layer) 238 | return x 239 | 240 | def _softmax(raw, input, dim=None, _stacklevel=3): 241 | # for F.softmax 242 | x=raw(input, dim=dim) 243 | if dim is None: 244 | dim=F._get_softmax_dim('softmax', input.dim(), _stacklevel) 245 | bottom_blobs=[log.blobs(input)] 246 | name = log.add_layer(name='softmax') 247 | log.add_blobs([x], name='softmax_blob') 248 | layer = caffe_net.Layer_param(name=name, type='Softmax', 249 | bottom=bottom_blobs, top=[log.blobs(x)]) 250 | layer.param.softmax_param.axis=dim 251 | log.cnet.add_layer(layer) 252 | return x 253 | 254 | def _batch_norm(raw,input, running_mean, running_var, weight=None, bias=None, 255 | training=False, momentum=0.1, eps=1e-5): 256 | # because the runing_mean and runing_var will be changed after the _batch_norm operation, we first save the parameters 257 | running_mean_clone=running_mean.clone() 258 | running_var_clone=running_var.clone() 259 | x = raw(input, running_mean, running_var, weight, bias, 260 | training, momentum, eps) 261 | bottom_blobs = [log.blobs(input)] 262 | layer_name1 = log.add_layer(name='batch_norm') 263 | top_blobs = log.add_blobs([x], name='batch_norm_blob') 264 | layer1 = caffe_net.Layer_param(name=layer_name1, type='BatchNorm', 265 | bottom=bottom_blobs, top=top_blobs) 266 | layer1.batch_norm_param(1, eps=eps) 267 | layer1.add_data(running_mean_clone.cpu().numpy(), running_var_clone.cpu().numpy(), np.array([1.0])) 268 | log.cnet.add_layer(layer1) 269 | layer_name2 = log.add_layer(name='bn_scale') 270 | layer2 = caffe_net.Layer_param(name=layer_name2, type='Scale', 271 | bottom=top_blobs, top=top_blobs)#top_blobs 272 | layer2.param.scale_param.bias_term = True 273 | layer2.add_data(weight.cpu().data.numpy(), bias.cpu().data.numpy()) 274 | log.cnet.add_layer(layer2) 275 | return x 276 | 277 | # ----- for Variable operations -------- 278 | 279 | def _view(input, *args): 280 | x=raw_view(input, *args) 281 | if not NET_INITTED: 282 | return x 283 | layer_name=log.add_layer(name='view') 284 | top_blobs=log.add_blobs([x],name='view_blob') 285 | layer=caffe_net.Layer_param(name=layer_name,type='Reshape', 286 | bottom=[log.blobs(input)],top=top_blobs) 287 | # TODO: reshpae added to nn_tools layer 288 | dims=list(args) 289 | dims[0]=0 # the first dim should be batch_size 290 | layer.param.reshape_param.shape.CopyFrom(caffe_net.pb.BlobShape(dim=dims)) 291 | log.cnet.add_layer(layer) 292 | return x 293 | 294 | def _add(input, *args): 295 | x = raw__add__(input, *args) 296 | if not NET_INITTED: 297 | return x 298 | layer_name = log.add_layer(name='add') 299 | top_blobs = log.add_blobs([x], name='add_blob') 300 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 301 | bottom=[log.blobs(input),log.blobs(args[0])], top=top_blobs) 302 | layer.param.eltwise_param.operation = 1 # sum is 1 303 | log.cnet.add_layer(layer) 304 | return x 305 | 306 | def _iadd(input, *args): 307 | x = raw__iadd__(input, *args) 308 | if not NET_INITTED: 309 | return x 310 | x=x.clone() 311 | layer_name = log.add_layer(name='add') 312 | top_blobs = log.add_blobs([x], name='add_blob') 313 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 314 | bottom=[log.blobs(input),log.blobs(args[0])], top=top_blobs) 315 | layer.param.eltwise_param.operation = 1 # sum is 1 316 | log.cnet.add_layer(layer) 317 | return x 318 | 319 | def _sub(input, *args): 320 | x = raw__sub__(input, *args) 321 | if not NET_INITTED: 322 | return x 323 | layer_name = log.add_layer(name='sub') 324 | top_blobs = log.add_blobs([x], name='sub_blob') 325 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 326 | bottom=[log.blobs(input),log.blobs(args[0])], top=top_blobs) 327 | layer.param.eltwise_param.operation = 1 # sum is 1 328 | layer.param.eltwise_param.coeff.extend([1.,-1.]) 329 | log.cnet.add_layer(layer) 330 | return x 331 | 332 | def _isub(input, *args): 333 | x = raw__isub__(input, *args) 334 | if not NET_INITTED: 335 | return x 336 | x=x.clone() 337 | layer_name = log.add_layer(name='sub') 338 | top_blobs = log.add_blobs([x], name='sub_blob') 339 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 340 | bottom=[log.blobs(input),log.blobs(args[0])], top=top_blobs) 341 | layer.param.eltwise_param.operation = 1 # sum is 1 342 | log.cnet.add_layer(layer) 343 | return x 344 | 345 | def _mul(input, *args): 346 | x = raw__sub__(input, *args) 347 | if not NET_INITTED: 348 | return x 349 | layer_name = log.add_layer(name='mul') 350 | top_blobs = log.add_blobs([x], name='mul_blob') 351 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 352 | bottom=[log.blobs(input), log.blobs(args[0])], top=top_blobs) 353 | layer.param.eltwise_param.operation = 0 # product is 1 354 | log.cnet.add_layer(layer) 355 | return x 356 | 357 | def _imul(input, *args): 358 | x = raw__isub__(input, *args) 359 | if not NET_INITTED: 360 | return x 361 | x = x.clone() 362 | layer_name = log.add_layer(name='mul') 363 | top_blobs = log.add_blobs([x], name='mul_blob') 364 | layer = caffe_net.Layer_param(name=layer_name, type='Eltwise', 365 | bottom=[log.blobs(input), log.blobs(args[0])], top=top_blobs) 366 | layer.param.eltwise_param.operation = 0 # product is 1 367 | layer.param.eltwise_param.coeff.extend([1., -1.]) 368 | log.cnet.add_layer(layer) 369 | return x 370 | 371 | class Rp(object): 372 | def __init__(self,raw,replace,**kwargs): 373 | # replace the raw function to replace function 374 | self.obj=replace 375 | self.raw=raw 376 | 377 | def __call__(self,*args,**kwargs): 378 | if not NET_INITTED: 379 | return self.raw(*args,**kwargs) 380 | out=self.obj(self.raw,*args,**kwargs) 381 | # if isinstance(out,Variable): 382 | # out=[out] 383 | return out 384 | 385 | 386 | F.conv2d=Rp(F.conv2d,_conv2d) 387 | F.linear=Rp(F.linear,_linear) 388 | F.max_pool2d=Rp(F.max_pool2d,_max_pool2d) 389 | F.avg_pool2d=Rp(F.avg_pool2d,_avg_pool2d) 390 | F.dropout=Rp(F.dropout,_dropout) 391 | F.threshold=Rp(F.threshold,_threshold) 392 | F.prelu=Rp(F.prelu,_prelu) 393 | F.batch_norm=Rp(F.batch_norm,_batch_norm) 394 | F.softmax=Rp(F.softmax,_softmax) 395 | 396 | torch.split=Rp(torch.split,_split) 397 | torch.max=Rp(torch.max,_max) 398 | torch.cat=Rp(torch.cat,_cat) 399 | 400 | 401 | # TODO: other types of the view function 402 | try: 403 | raw_view=Variable.view 404 | Variable.view=_view 405 | raw__add__=Variable.__add__ 406 | Variable.__add__=_add 407 | raw__iadd__=Variable.__iadd__ 408 | Variable.__iadd__=_iadd 409 | raw__sub__=Variable.__sub__ 410 | Variable.__sub__=_sub 411 | raw__isub__=Variable.__isub__ 412 | Variable.__isub__=_isub 413 | raw__mul__ = Variable.__mul__ 414 | Variable.__mul__ = _mul 415 | raw__imul__ = Variable.__imul__ 416 | Variable.__imul__ = _imul 417 | except: 418 | # for new version 0.4.0 419 | for t in [torch.Tensor]: 420 | raw_view = t.view 421 | t.view = _view 422 | raw__add__ = t.__add__ 423 | t.__add__ = _add 424 | raw__iadd__ = t.__iadd__ 425 | t.__iadd__ = _iadd 426 | raw__sub__ = t.__sub__ 427 | t.__sub__ = _sub 428 | raw__isub__ = t.__isub__ 429 | t.__isub__ = _isub 430 | raw__mul__ = t.__mul__ 431 | t.__mul__=_mul 432 | raw__imul__ = t.__imul__ 433 | t.__imul__ = _imul 434 | 435 | 436 | def trans_net(net,input_var,name='NoNamePytorchModel'): 437 | print('Starting Transform, This will take a while') 438 | log.init([input_var]) 439 | log.cnet.net.name=name 440 | log.cnet.net.input.extend([log.blobs(input_var)]) 441 | log.cnet.net.input_dim.extend(input_var.size()) 442 | global NET_INITTED 443 | NET_INITTED=True 444 | out = net.forward(input_var) 445 | print('Transform Completed') 446 | 447 | def save_prototxt(save_name): 448 | log.cnet.save_prototxt(save_name) 449 | 450 | def save_caffemodel(save_name): 451 | log.cnet.save(save_name) 452 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/wdd0225/pytorch2caffe/479a18308903e502f74dce6d2b4c8dede2ef62ee/utils/__init__.py -------------------------------------------------------------------------------- /utils/functions.py: -------------------------------------------------------------------------------- 1 | from collections import defaultdict 2 | import numpy as np 3 | import torch 4 | from sklearn.metrics import average_precision_score 5 | 6 | 7 | def _unique_sample(ids_dict, num): 8 | mask = np.zeros(num, dtype=np.bool) 9 | for _, indices in ids_dict.items(): 10 | i = np.random.choice(indices) 11 | mask[i] = True 12 | return mask 13 | 14 | 15 | def cmc(distmat, query_ids=None, gallery_ids=None, 16 | query_cams=None, gallery_cams=None, topk=100, 17 | separate_camera_set=False, 18 | single_gallery_shot=False, 19 | first_match_break=False): 20 | m, n = distmat.shape 21 | # Fill up default values 22 | if query_ids is None: 23 | query_ids = np.arange(m) 24 | if gallery_ids is None: 25 | gallery_ids = np.arange(n) 26 | if query_cams is None: 27 | query_cams = np.zeros(m).astype(np.int32) 28 | if gallery_cams is None: 29 | gallery_cams = np.ones(n).astype(np.int32) 30 | # Ensure numpy array 31 | query_ids = np.asarray(query_ids) 32 | gallery_ids = np.asarray(gallery_ids) 33 | query_cams = np.asarray(query_cams) 34 | gallery_cams = np.asarray(gallery_cams) 35 | # Sort and find correct matches 36 | indices = np.argsort(distmat, axis=1) 37 | matches = (gallery_ids[indices] == query_ids[:, np.newaxis]) 38 | # Compute CMC for each query 39 | ret = np.zeros(topk) 40 | num_valid_queries = 0 41 | for i in range(m): 42 | # Filter out the same id and same camera 43 | valid = ((gallery_ids[indices[i]] != query_ids[i]) | 44 | (gallery_cams[indices[i]] != query_cams[i])) 45 | if separate_camera_set: 46 | # Filter out samples from same camera 47 | valid &= (gallery_cams[indices[i]] != query_cams[i]) 48 | if not np.any(matches[i, valid]): 49 | continue 50 | if single_gallery_shot: 51 | repeat = 10 52 | gids = gallery_ids[indices[i][valid]] 53 | inds = np.where(valid)[0] 54 | ids_dict = defaultdict(list) 55 | for j, x in zip(inds, gids): 56 | ids_dict[x].append(j) 57 | else: 58 | repeat = 1 59 | for _ in range(repeat): 60 | if single_gallery_shot: 61 | # Randomly choose one instance for each id 62 | sampled = (valid & _unique_sample(ids_dict, len(valid))) 63 | index = np.nonzero(matches[i, sampled])[0] 64 | else: 65 | index = np.nonzero(matches[i, valid])[0] 66 | delta = 1. / (len(index) * repeat) 67 | for j, k in enumerate(index): 68 | if k - j >= topk: 69 | break 70 | if first_match_break: 71 | ret[k - j] += 1 72 | break 73 | ret[k - j] += delta 74 | num_valid_queries += 1 75 | if num_valid_queries == 0: 76 | raise RuntimeError("No valid query") 77 | return ret.cumsum() / num_valid_queries 78 | 79 | 80 | def mean_ap(distmat, query_ids=None, gallery_ids=None, 81 | query_cams=None, gallery_cams=None): 82 | m, n = distmat.shape 83 | # Fill up default values 84 | if query_ids is None: 85 | query_ids = np.arange(m) 86 | if gallery_ids is None: 87 | gallery_ids = np.arange(n) 88 | if query_cams is None: 89 | query_cams = np.zeros(m).astype(np.int32) 90 | if gallery_cams is None: 91 | gallery_cams = np.ones(n).astype(np.int32) 92 | # Ensure numpy array 93 | query_ids = np.asarray(query_ids) 94 | gallery_ids = np.asarray(gallery_ids) 95 | query_cams = np.asarray(query_cams) 96 | gallery_cams = np.asarray(gallery_cams) 97 | # Sort and find correct matches 98 | indices = np.argsort(distmat, axis=1) 99 | matches = (gallery_ids[indices] == query_ids[:, np.newaxis]) 100 | # Compute AP for each query 101 | aps = [] 102 | for i in range(m): 103 | # Filter out the same id and same camera 104 | valid = ((gallery_ids[indices[i]] != query_ids[i]) | 105 | (gallery_cams[indices[i]] != query_cams[i])) 106 | y_true = matches[i, valid] 107 | y_score = -distmat[i][indices[i]][valid] 108 | if not np.any(y_true): 109 | continue 110 | aps.append(average_precision_score(y_true, y_score)) 111 | if len(aps) == 0: 112 | raise RuntimeError("No valid query") 113 | return np.mean(aps) 114 | -------------------------------------------------------------------------------- /utils/random_erasing.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | from torchvision.transforms import * 4 | 5 | from PIL import Image 6 | import random 7 | import math 8 | import numpy as np 9 | import torch 10 | 11 | class RandomErasing(object): 12 | """ Randomly selects a rectangle region in an image and erases its pixels. 13 | 'Random Erasing Data Augmentation' by Zhong et al. 14 | See https://arxiv.org/pdf/1708.04896.pdf 15 | Args: 16 | probability: The probability that the Random Erasing operation will be performed. 17 | sl: Minimum proportion of erased area against input image. 18 | sh: Maximum proportion of erased area against input image. 19 | r1: Minimum aspect ratio of erased area. 20 | mean: Erasing value. 21 | """ 22 | 23 | def __init__(self, probability = 0.5, sl = 0.02, sh = 0.4, r1 = 0.3, mean=[0.4914, 0.4822, 0.4465]): 24 | self.probability = probability 25 | self.mean = mean 26 | self.sl = sl 27 | self.sh = sh 28 | self.r1 = r1 29 | 30 | def __call__(self, img): 31 | 32 | if random.uniform(0, 1) > self.probability: 33 | return img 34 | 35 | for attempt in range(100): 36 | area = img.size()[1] * img.size()[2] 37 | 38 | target_area = random.uniform(self.sl, self.sh) * area 39 | aspect_ratio = random.uniform(self.r1, 1/self.r1) 40 | 41 | h = int(round(math.sqrt(target_area * aspect_ratio))) 42 | w = int(round(math.sqrt(target_area / aspect_ratio))) 43 | 44 | if w < img.size()[2] and h < img.size()[1]: 45 | x1 = random.randint(0, img.size()[1] - h) 46 | y1 = random.randint(0, img.size()[2] - w) 47 | if img.size()[0] == 3: 48 | img[0, x1:x1+h, y1:y1+w] = self.mean[0] 49 | img[1, x1:x1+h, y1:y1+w] = self.mean[1] 50 | img[2, x1:x1+h, y1:y1+w] = self.mean[2] 51 | else: 52 | img[0, x1:x1+h, y1:y1+w] = self.mean[0] 53 | return img 54 | 55 | return img 56 | -------------------------------------------------------------------------------- /utils/re_ranking.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python2/python3 2 | # -*- coding: utf-8 -*- 3 | """ 4 | Created on Mon Jun 26 14:46:56 2017 5 | @author: luohao 6 | Modified by Houjing Huang, 2017-12-22. 7 | - This version accepts distance matrix instead of raw features. 8 | - The difference of `/` division between python 2 and 3 is handled. 9 | - numpy.float16 is replaced by numpy.float32 for numerical precision. 10 | 11 | Modified by Zhedong Zheng, 2018-1-12. 12 | - replace sort with topK, which save about 30s. 13 | """ 14 | 15 | """ 16 | CVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-identification with k-reciprocal Encoding[J]. 2017. 17 | url:http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhong_Re-Ranking_Person_Re-Identification_CVPR_2017_paper.pdf 18 | Matlab version: https://github.com/zhunzhong07/person-re-ranking 19 | """ 20 | 21 | """ 22 | API 23 | q_g_dist: query-gallery distance matrix, numpy array, shape [num_query, num_gallery] 24 | q_q_dist: query-query distance matrix, numpy array, shape [num_query, num_query] 25 | g_g_dist: gallery-gallery distance matrix, numpy array, shape [num_gallery, num_gallery] 26 | k1, k2, lambda_value: parameters, the original paper is (k1=20, k2=6, lambda_value=0.3) 27 | Returns: 28 | final_dist: re-ranked distance, numpy array, shape [num_query, num_gallery] 29 | """ 30 | 31 | 32 | import numpy as np 33 | 34 | def k_reciprocal_neigh( initial_rank, i, k1): 35 | forward_k_neigh_index = initial_rank[i,:k1+1] 36 | backward_k_neigh_index = initial_rank[forward_k_neigh_index,:k1+1] 37 | fi = np.where(backward_k_neigh_index==i)[0] 38 | return forward_k_neigh_index[fi] 39 | 40 | def re_ranking(q_g_dist, q_q_dist, g_g_dist, k1=20, k2=6, lambda_value=0.3): 41 | # The following naming, e.g. gallery_num, is different from outer scope. 42 | # Don't care about it. 43 | original_dist = np.concatenate( 44 | [np.concatenate([q_q_dist, q_g_dist], axis=1), 45 | np.concatenate([q_g_dist.T, g_g_dist], axis=1)], 46 | axis=0) 47 | original_dist = 2. - 2 * original_dist #np.power(original_dist, 2).astype(np.float32) 48 | original_dist = np.transpose(1. * original_dist/np.max(original_dist,axis = 0)) 49 | V = np.zeros_like(original_dist).astype(np.float32) 50 | #initial_rank = np.argsort(original_dist).astype(np.int32) 51 | # top K1+1 52 | initial_rank = np.argpartition( original_dist, range(1,k1+1) ) 53 | 54 | query_num = q_g_dist.shape[0] 55 | all_num = original_dist.shape[0] 56 | 57 | for i in range(all_num): 58 | # k-reciprocal neighbors 59 | k_reciprocal_index = k_reciprocal_neigh( initial_rank, i, k1) 60 | k_reciprocal_expansion_index = k_reciprocal_index 61 | for j in range(len(k_reciprocal_index)): 62 | candidate = k_reciprocal_index[j] 63 | candidate_k_reciprocal_index = k_reciprocal_neigh( initial_rank, candidate, int(np.around(k1/2))) 64 | if len(np.intersect1d(candidate_k_reciprocal_index,k_reciprocal_index))> 2./3*len(candidate_k_reciprocal_index): 65 | k_reciprocal_expansion_index = np.append(k_reciprocal_expansion_index,candidate_k_reciprocal_index) 66 | 67 | k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index) 68 | weight = np.exp(-original_dist[i,k_reciprocal_expansion_index]) 69 | V[i,k_reciprocal_expansion_index] = 1.*weight/np.sum(weight) 70 | 71 | original_dist = original_dist[:query_num,] 72 | if k2 != 1: 73 | V_qe = np.zeros_like(V,dtype=np.float32) 74 | for i in range(all_num): 75 | V_qe[i,:] = np.mean(V[initial_rank[i,:k2],:],axis=0) 76 | V = V_qe 77 | del V_qe 78 | del initial_rank 79 | invIndex = [] 80 | for i in range(all_num): 81 | invIndex.append(np.where(V[:,i] != 0)[0]) 82 | 83 | jaccard_dist = np.zeros_like(original_dist,dtype = np.float32) 84 | 85 | for i in range(query_num): 86 | temp_min = np.zeros(shape=[1,all_num],dtype=np.float32) 87 | indNonZero = np.where(V[i,:] != 0)[0] 88 | indImages = [] 89 | indImages = [invIndex[ind] for ind in indNonZero] 90 | for j in range(len(indNonZero)): 91 | temp_min[0,indImages[j]] = temp_min[0,indImages[j]]+ np.minimum(V[i,indNonZero[j]],V[indImages[j],indNonZero[j]]) 92 | jaccard_dist[i] = 1-temp_min/(2.-temp_min) 93 | 94 | final_dist = jaccard_dist*(1-lambda_value) + original_dist*lambda_value 95 | del original_dist 96 | del V 97 | del jaccard_dist 98 | final_dist = final_dist[:query_num,query_num:] 99 | return final_dist 100 | -------------------------------------------------------------------------------- /utils/utility.py: -------------------------------------------------------------------------------- 1 | import os 2 | import datetime 3 | 4 | import matplotlib 5 | matplotlib.use('Agg') 6 | import matplotlib.pyplot as plt 7 | 8 | import numpy as np 9 | import scipy.misc as misc 10 | 11 | import torch 12 | import torch.optim as optim 13 | import torch.optim.lr_scheduler as lrs 14 | 15 | class checkpoint(): 16 | def __init__(self, args): 17 | self.args = args 18 | self.log = torch.Tensor() 19 | now = datetime.datetime.now().strftime('%Y-%m-%d-%H:%M:%S') 20 | 21 | if args.load == '': 22 | if args.save == '': args.save = now 23 | self.dir = 'experiment/' + args.save 24 | else: 25 | self.dir = 'experiment/' + args.load 26 | if not os.path.exists(self.dir): 27 | args.load = '' 28 | else: 29 | self.log = torch.load(self.dir + '/map_log.pt') 30 | print('Continue from epoch {}...'.format(len(self.log)*args.test_every)) 31 | 32 | if args.reset: 33 | os.system('rm -rf ' + self.dir) 34 | args.load = '' 35 | 36 | def _make_dir(path): 37 | if not os.path.exists(path): os.makedirs(path) 38 | 39 | _make_dir(self.dir) 40 | _make_dir(self.dir + '/model') 41 | _make_dir(self.dir + '/results') 42 | 43 | open_type = 'a' if os.path.exists(self.dir + '/log.txt') else 'w' 44 | self.log_file = open(self.dir + '/log.txt', open_type) 45 | with open(self.dir + '/config.txt', open_type) as f: 46 | f.write(now + '\n\n') 47 | for arg in vars(args): 48 | f.write('{}: {}\n'.format(arg, getattr(args, arg))) 49 | f.write('\n') 50 | 51 | def save(self, trainer, epoch, is_best=False): 52 | trainer.model.save(self.dir, epoch, is_best=is_best) 53 | trainer.loss.save(self.dir) 54 | trainer.loss.plot_loss(self.dir, epoch) 55 | 56 | self.plot_map_rank(epoch) 57 | torch.save(self.log, os.path.join(self.dir, 'map_log.pt')) 58 | torch.save( 59 | trainer.optimizer.state_dict(), 60 | os.path.join(self.dir, 'optimizer.pt') 61 | ) 62 | 63 | def add_log(self, log): 64 | self.log = torch.cat([self.log, log]) 65 | 66 | def write_log(self, log, refresh=False, end='\n'): 67 | print(log, end=end) 68 | if end != '': 69 | self.log_file.write(log + end) 70 | if refresh: 71 | self.log_file.close() 72 | self.log_file = open(self.dir + '/log.txt', 'a') 73 | 74 | def done(self): 75 | self.log_file.close() 76 | 77 | def plot_map_rank(self, epoch): 78 | axis = np.linspace(1, epoch, self.log.size(0)) 79 | label = 'Reid on {}'.format(self.args.data_test) 80 | labels = ['mAP','rank1','rank3','rank5','rank10'] 81 | fig = plt.figure() 82 | plt.title(label) 83 | for i in range(len(labels)): 84 | plt.plot(axis, self.log[:, i].numpy(), label=labels[i]) 85 | 86 | plt.legend() 87 | plt.xlabel('Epochs') 88 | plt.ylabel('mAP/rank') 89 | plt.grid(True) 90 | plt.savefig('{}/test_{}.jpg'.format(self.dir, self.args.data_test)) 91 | plt.close(fig) 92 | 93 | def save_results(self, filename, save_list, scale): 94 | pass 95 | 96 | def make_optimizer(args, model): 97 | trainable = filter(lambda x: x.requires_grad, model.parameters()) 98 | 99 | if args.optimizer == 'SGD': 100 | optimizer_function = optim.SGD 101 | kwargs = { 102 | 'momentum': args.momentum, 103 | 'dampening': args.dampening, 104 | 'nesterov': args.nesterov 105 | } 106 | elif args.optimizer == 'ADAM': 107 | optimizer_function = optim.Adam 108 | kwargs = { 109 | 'betas': (args.beta1, args.beta2), 110 | 'eps': args.epsilon, 111 | 'amsgrad': args.amsgrad 112 | } 113 | elif args.optimizer == 'ADAMAX': 114 | optimizer_function = optim.Adamax 115 | kwargs = { 116 | 'betas': (args.beta1, args.beta2), 117 | 'eps': args.epsilon 118 | } 119 | elif args.optimizer == 'RMSprop': 120 | optimizer_function = optim.RMSprop 121 | kwargs = { 122 | 'eps': args.epsilon, 123 | 'momentum': args.momentum 124 | } 125 | else: 126 | raise Exception() 127 | 128 | kwargs['lr'] = args.lr 129 | kwargs['weight_decay'] = args.weight_decay 130 | 131 | return optimizer_function(trainable, **kwargs) 132 | 133 | def make_scheduler(args, optimizer): 134 | if args.decay_type == 'step': 135 | scheduler = lrs.StepLR( 136 | optimizer, 137 | step_size=args.lr_decay, 138 | gamma=args.gamma 139 | ) 140 | elif args.decay_type.find('step') >= 0: 141 | milestones = args.decay_type.split('_') 142 | milestones.pop(0) 143 | milestones = list(map(lambda x: int(x), milestones)) 144 | scheduler = lrs.MultiStepLR( 145 | optimizer, 146 | milestones=milestones, 147 | gamma=args.gamma 148 | ) 149 | 150 | return scheduler 151 | 152 | --------------------------------------------------------------------------------