├── .gitignore ├── README.md ├── convert_tf_to_pt ├── convert_params_tf_pytorch.py ├── model.py ├── model_caffe.py ├── original_tf │ ├── __init__.py │ ├── efficientnet_builder.py │ ├── efficientnet_model.py │ ├── eval_ckpt_main.py │ ├── preprocessing.py │ └── utils.py ├── prototxt_basic.py ├── pytorch2caffe.py ├── pytorch2caffe_model.py ├── test.jpg ├── test_caffe.py ├── test_pytorch.py └── utils.py └── pretrained_tensorflow └── download.sh /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *pyc 3 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ### build 2 | * pip install tensorflow torch torchvision 3 | * pip3 install tensorflow torch torchvision 4 | * git clone https://github.com/xxworkspace/caffe.git 5 | * Then build caffe width cpu-only 6 | 7 | ### TensorFlow to PyTorch Conversion 8 | * Download tf-model 9 | * cd pretrained_tensorflow 10 | * bash download.sh efficientnet-b* 11 | * Convert tf-model to pytorch-model 12 | * cd convert_tf_pt 13 | * python3 convert_params_tf_pytorch.py --model_name efficientnet-b0 --tf_checkpoint ../pretrained_tensorflow/efficientnet-b0/ --output_file ../pretrained_pytorch/efficientnet-b0.pth 14 | * Using python3 convert_params_tf_pytorch.py -h 15 | * Test 16 | * python test_pytorch.py efficientnet-b0* 17 | 18 | ### PyTorch to Caffe 19 | * Convert pytorch-model to caffe .prototxt 20 | * python pytorch2caffe.py efficientnet-b* 21 | * Convert pytorch-model to caffe .caffemodel 22 | * python pytorch2caffe_model.py efficientnet-b* 23 | * Test 24 | * pytorch test_caffe.py efficientnet-b* 25 | 26 | ### data process RGB 27 | mean = [0.485, 0.456, 0.406] 28 | std = [0.229, 0.224, 0.225] 29 | -------------------------------------------------------------------------------- /convert_tf_to_pt/convert_params_tf_pytorch.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | import torch 4 | 5 | ''' 6 | def load_param(checkpoint_file, conversion_table, model_name): 7 | """ 8 | Load parameters according to conversion_table. 9 | 10 | Args: 11 | checkpoint_file (string): pretrained checkpoint model file in tensorflow 12 | conversion_table (dict): { pytorch tensor in a model : checkpoint variable name } 13 | """ 14 | dict = {} 15 | for pyt_param, tf_param_name in conversion_table.items(): 16 | param_name = tf_param_name 17 | tf_param_name = str(model_name) + '/' + tf_param_name 18 | tf_param = tf.train.load_variable(checkpoint_file, tf_param_name) 19 | if 'conv' in tf_param_name and 'kernel' in tf_param_name: 20 | tf_param = np.transpose(tf_param, (3, 2, 0, 1)) 21 | if 'depthwise' in tf_param_name: 22 | tf_param = np.transpose(tf_param, (1, 0, 2, 3)) 23 | elif tf_param_name.endswith('kernel'): # for weight(kernel), we should do transpose 24 | tf_param = np.transpose(tf_param) 25 | assert pyt_param.size() == tf_param.shape, \ 26 | 'Dim Mismatch: %s vs %s ; %s' % (tuple(pyt_param.size()), tf_param.shape, tf_param_name) 27 | print(tf_param_name," --> ",type(pyt_param)) 28 | #print(type(pyt_param)) 29 | pyt_param = torch.from_numpy(tf_param) 30 | dict[param_name] = tf_param.copy() 31 | #print(pyt_param[0]) 32 | #[pyt_param] = tf_param_name 33 | #print(tf_param_name,tf_param) 34 | #print(pyt_param) 35 | #dict[tf_param_name] = tf_param.copy() 36 | #print(tf_param) 37 | #if "bias" in tf_param_name: 38 | #print(pyt_param) 39 | #print(conversion_table) 40 | return dict 41 | ''' 42 | 43 | def pytorch_param_from_tf_checkpoint(checkpoint_file, conversion_table, params_shape_table, model_name): 44 | """ 45 | Load parameters according to conversion_table. 46 | 47 | Args: 48 | checkpoint_file (string): pretrained checkpoint model file in tensorflow 49 | conversion_table (dict): { pytorch tensor in a model : checkpoint variable name } 50 | """ 51 | dict = {} 52 | params_shape_table_new = {} 53 | for k,v in params_shape_table.items(): 54 | params_shape_table_new[v] = k 55 | 56 | for pyt_param_name, tf_param_name in conversion_table.items(): 57 | tf_param_name_ = tf_param_name 58 | tf_param_name = str(model_name) + '/' + tf_param_name 59 | tf_param = tf.train.load_variable(checkpoint_file, tf_param_name) 60 | if 'conv' in tf_param_name and 'kernel' in tf_param_name: 61 | tf_param = np.transpose(tf_param, (3, 2, 0, 1)) 62 | if 'depthwise' in tf_param_name: 63 | tf_param = np.transpose(tf_param, (1, 0, 2, 3)) 64 | elif tf_param_name.endswith('kernel'): 65 | tf_param = np.transpose(tf_param) 66 | assert params_shape_table_new[tf_param_name_].size() == tf_param.shape, \ 67 | 'Dim Mismatch: %s vs %s ; %s' % (tuple(params_shape_table_new[tf_param_name_].size()), tf_param.shape, tf_param_name) 68 | print(tf_param_name," --> ",pyt_param_name) 69 | dict[pyt_param_name] = tf_param.copy() 70 | 71 | return dict 72 | 73 | def pytorch_tf_name_dict(model): 74 | # This will store the enire conversion table 75 | conversion_table = {} 76 | merge = lambda dict1, dict2: {**dict1, **dict2} 77 | 78 | # All the weights not in the conv blocks 79 | conversion_table_for_weights_outside_blocks = { 80 | "_conv_stem.weight": 'stem/conv2d/kernel', # [3, 3, 3, 32]), 81 | "_bn0.bias": 'stem/tpu_batch_normalization/beta', # [32]), 82 | "_bn0.weight": 'stem/tpu_batch_normalization/gamma', # [32]), 83 | "_bn0.running_mean": 'stem/tpu_batch_normalization/moving_mean', # [32]), 84 | "_bn0.running_var": 'stem/tpu_batch_normalization/moving_variance', # [32]), 85 | "_conv_head.weight": 'head/conv2d/kernel', # [1, 1, 320, 1280]), 86 | "_bn1.bias": 'head/tpu_batch_normalization/beta', # [1280]), 87 | "_bn1.weight": 'head/tpu_batch_normalization/gamma', # [1280]), 88 | "_bn1.running_mean": 'head/tpu_batch_normalization/moving_mean', # [32]), 89 | "_bn1.running_var": 'head/tpu_batch_normalization/moving_variance', # [32]), 90 | "_fc.bias": 'head/dense/bias', # [1000]), 91 | "_fc.weight": 'head/dense/kernel', # [1280, 1000]), 92 | } 93 | conversion_table = merge(conversion_table, conversion_table_for_weights_outside_blocks) 94 | 95 | # The first conv block is special because it does not have _expand_conv 96 | conversion_table_for_first_block = { 97 | "_blocks.0._project_conv.weight": 'blocks_0/conv2d/kernel', # 1, 1, 32, 16]), 98 | "_blocks.0._depthwise_conv.weight": 'blocks_0/depthwise_conv2d/depthwise_kernel', # [3, 3, 32, 1]), 99 | "_blocks.0._se_reduce.bias": 'blocks_0/se/conv2d/bias', # , [8]), 100 | "_blocks.0._se_reduce.weight": 'blocks_0/se/conv2d/kernel', # , [1, 1, 32, 8]), 101 | "_blocks.0._se_expand.bias": 'blocks_0/se/conv2d_1/bias', # , [32]), 102 | "_blocks.0._se_expand.weight": 'blocks_0/se/conv2d_1/kernel', # , [1, 1, 8, 32]), 103 | "_blocks.0._bn1.bias": 'blocks_0/tpu_batch_normalization/beta', # [32]), 104 | "_blocks.0._bn1.weight": 'blocks_0/tpu_batch_normalization/gamma', # [32]), 105 | "_blocks.0._bn1.running_mean": 'blocks_0/tpu_batch_normalization/moving_mean', 106 | "_blocks.0._bn1.running_var": 'blocks_0/tpu_batch_normalization/moving_variance', 107 | "_blocks.0._bn2.bias": 'blocks_0/tpu_batch_normalization_1/beta', # [16]), 108 | "_blocks.0._bn2.weight": 'blocks_0/tpu_batch_normalization_1/gamma', # [16]), 109 | "_blocks.0._bn2.running_mean": 'blocks_0/tpu_batch_normalization_1/moving_mean', 110 | "_blocks.0._bn2.running_var": 'blocks_0/tpu_batch_normalization_1/moving_variance', 111 | } 112 | conversion_table = merge(conversion_table, conversion_table_for_first_block) 113 | 114 | # Conv blocks 115 | for i in range(len(model._blocks)): 116 | 117 | is_first_block = '_expand_conv.weight' not in [n for n, p in model._blocks[i].named_parameters()] 118 | 119 | if is_first_block: 120 | conversion_table_block = { 121 | "_blocks." + str(i) +"._project_conv.weight": 'blocks_' + str(i) + '/conv2d/kernel', # 1, 1, 32, 16]), 122 | "_blocks." + str(i) +"._depthwise_conv.weight": 'blocks_' + str(i) + '/depthwise_conv2d/depthwise_kernel', 123 | # [3, 3, 32, 1]), 124 | "_blocks." + str(i) +"._se_reduce.bias": 'blocks_' + str(i) + '/se/conv2d/bias', # , [8]), 125 | "_blocks." + str(i) +"._se_reduce.weight": 'blocks_' + str(i) + '/se/conv2d/kernel', # , [1, 1, 32, 8]), 126 | "_blocks." + str(i) +"._se_expand.bias": 'blocks_' + str(i) + '/se/conv2d_1/bias', # , [32]), 127 | "_blocks." + str(i) +"._se_expand.weight": 'blocks_' + str(i) + '/se/conv2d_1/kernel', # , [1, 1, 8, 32]), 128 | "_blocks." + str(i) +"._bn1.bias": 'blocks_' + str(i) + '/tpu_batch_normalization/beta', # [32]), 129 | "_blocks." + str(i) +"._bn1.weight": 'blocks_' + str(i) + '/tpu_batch_normalization/gamma', # [32]), 130 | "_blocks." + str(i) +"._bn1.running_mean": 'blocks_' + str(i) + '/tpu_batch_normalization/moving_mean', 131 | "_blocks." + str(i) +"._bn1.running_var": 'blocks_' + str(i) + '/tpu_batch_normalization/moving_variance', 132 | "_blocks." + str(i) +"._bn2.bias": 'blocks_' + str(i) + '/tpu_batch_normalization_1/beta', # [16]), 133 | "_blocks." + str(i) +"._bn2.weight": 'blocks_' + str(i) + '/tpu_batch_normalization_1/gamma', # [16]), 134 | "_blocks." + str(i) +"._bn2.running_mean": 'blocks_' + str(i) + '/tpu_batch_normalization_1/moving_mean', 135 | "_blocks." + str(i) +"._bn2.running_var": 'blocks_' + str(i) + '/tpu_batch_normalization_1/moving_variance', 136 | } 137 | 138 | else: 139 | conversion_table_block = { 140 | "_blocks." + str(i) +"._expand_conv.weight": 'blocks_' + str(i) + '/conv2d/kernel', 141 | "_blocks." + str(i) +"._project_conv.weight": 'blocks_' + str(i) + '/conv2d_1/kernel', 142 | "_blocks." + str(i) +"._depthwise_conv.weight": 'blocks_' + str(i) + '/depthwise_conv2d/depthwise_kernel', 143 | "_blocks." + str(i) +"._se_reduce.bias": 'blocks_' + str(i) + '/se/conv2d/bias', 144 | "_blocks." + str(i) +"._se_reduce.weight": 'blocks_' + str(i) + '/se/conv2d/kernel', 145 | "_blocks." + str(i) +"._se_expand.bias": 'blocks_' + str(i) + '/se/conv2d_1/bias', 146 | "_blocks." + str(i) +"._se_expand.weight": 'blocks_' + str(i) + '/se/conv2d_1/kernel', 147 | "_blocks." + str(i) +"._bn0.bias": 'blocks_' + str(i) + '/tpu_batch_normalization/beta', 148 | "_blocks." + str(i) +"._bn0.weight": 'blocks_' + str(i) + '/tpu_batch_normalization/gamma', 149 | "_blocks." + str(i) +"._bn0.running_mean": 'blocks_' + str(i) + '/tpu_batch_normalization/moving_mean', 150 | "_blocks." + str(i) +"._bn0.running_var": 'blocks_' + str(i) + '/tpu_batch_normalization/moving_variance', 151 | "_blocks." + str(i) +"._bn1.bias": 'blocks_' + str(i) + '/tpu_batch_normalization_1/beta', 152 | "_blocks." + str(i) +"._bn1.weight": 'blocks_' + str(i) + '/tpu_batch_normalization_1/gamma', 153 | "_blocks." + str(i) +"._bn1.running_mean": 'blocks_' + str(i) + '/tpu_batch_normalization_1/moving_mean', 154 | "_blocks." + str(i) +"._bn1.running_var": 'blocks_' + str(i) + '/tpu_batch_normalization_1/moving_variance', 155 | "_blocks." + str(i) +"._bn2.bias": 'blocks_' + str(i) + '/tpu_batch_normalization_2/beta', 156 | "_blocks." + str(i) +"._bn2.weight": 'blocks_' + str(i) + '/tpu_batch_normalization_2/gamma', 157 | "_blocks." + str(i) +"._bn2.running_mean": 'blocks_' + str(i) + '/tpu_batch_normalization_2/moving_mean', 158 | "_blocks." + str(i) +"._bn2.running_var": 'blocks_' + str(i) + '/tpu_batch_normalization_2/moving_variance', 159 | } 160 | 161 | conversion_table = merge(conversion_table, conversion_table_block) 162 | #return load_param_from_tf(checkpoint_file,conversion_table,model_name) 163 | return conversion_table 164 | 165 | def pytorch_tf_params_dict(model): 166 | """ 167 | 168 | """ 169 | # This will store the enire conversion table 170 | conversion_table = {} 171 | merge = lambda dict1, dict2: {**dict1, **dict2} 172 | 173 | # All the weights not in the conv blocks 174 | conversion_table_for_weights_outside_blocks = { 175 | model._conv_stem.weight: 'stem/conv2d/kernel', # [3, 3, 3, 32]), 176 | model._bn0.bias: 'stem/tpu_batch_normalization/beta', # [32]), 177 | model._bn0.weight: 'stem/tpu_batch_normalization/gamma', # [32]), 178 | model._bn0.running_mean: 'stem/tpu_batch_normalization/moving_mean', # [32]), 179 | model._bn0.running_var: 'stem/tpu_batch_normalization/moving_variance', # [32]), 180 | model._conv_head.weight: 'head/conv2d/kernel', # [1, 1, 320, 1280]), 181 | model._bn1.bias: 'head/tpu_batch_normalization/beta', # [1280]), 182 | model._bn1.weight: 'head/tpu_batch_normalization/gamma', # [1280]), 183 | model._bn1.running_mean: 'head/tpu_batch_normalization/moving_mean', # [32]), 184 | model._bn1.running_var: 'head/tpu_batch_normalization/moving_variance', # [32]), 185 | model._fc.bias: 'head/dense/bias', # [1000]), 186 | model._fc.weight: 'head/dense/kernel', # [1280, 1000]), 187 | } 188 | conversion_table = merge(conversion_table, conversion_table_for_weights_outside_blocks) 189 | 190 | # The first conv block is special because it does not have _expand_conv 191 | conversion_table_for_first_block = { 192 | model._blocks[0]._project_conv.weight: 'blocks_0/conv2d/kernel', # 1, 1, 32, 16]), 193 | model._blocks[0]._depthwise_conv.weight: 'blocks_0/depthwise_conv2d/depthwise_kernel', # [3, 3, 32, 1]), 194 | model._blocks[0]._se_reduce.bias: 'blocks_0/se/conv2d/bias', # , [8]), 195 | model._blocks[0]._se_reduce.weight: 'blocks_0/se/conv2d/kernel', # , [1, 1, 32, 8]), 196 | model._blocks[0]._se_expand.bias: 'blocks_0/se/conv2d_1/bias', # , [32]), 197 | model._blocks[0]._se_expand.weight: 'blocks_0/se/conv2d_1/kernel', # , [1, 1, 8, 32]), 198 | model._blocks[0]._bn1.bias: 'blocks_0/tpu_batch_normalization/beta', # [32]), 199 | model._blocks[0]._bn1.weight: 'blocks_0/tpu_batch_normalization/gamma', # [32]), 200 | model._blocks[0]._bn1.running_mean: 'blocks_0/tpu_batch_normalization/moving_mean', 201 | model._blocks[0]._bn1.running_var: 'blocks_0/tpu_batch_normalization/moving_variance', 202 | model._blocks[0]._bn2.bias: 'blocks_0/tpu_batch_normalization_1/beta', # [16]), 203 | model._blocks[0]._bn2.weight: 'blocks_0/tpu_batch_normalization_1/gamma', # [16]), 204 | model._blocks[0]._bn2.running_mean: 'blocks_0/tpu_batch_normalization_1/moving_mean', 205 | model._blocks[0]._bn2.running_var: 'blocks_0/tpu_batch_normalization_1/moving_variance', 206 | } 207 | conversion_table = merge(conversion_table, conversion_table_for_first_block) 208 | 209 | # Conv blocks 210 | for i in range(len(model._blocks)): 211 | 212 | is_first_block = '_expand_conv.weight' not in [n for n, p in model._blocks[i].named_parameters()] 213 | 214 | if is_first_block: 215 | conversion_table_block = { 216 | model._blocks[i]._project_conv.weight: 'blocks_' + str(i) + '/conv2d/kernel', # 1, 1, 32, 16]), 217 | model._blocks[i]._depthwise_conv.weight: 'blocks_' + str(i) + '/depthwise_conv2d/depthwise_kernel', 218 | # [3, 3, 32, 1]), 219 | model._blocks[i]._se_reduce.bias: 'blocks_' + str(i) + '/se/conv2d/bias', # , [8]), 220 | model._blocks[i]._se_reduce.weight: 'blocks_' + str(i) + '/se/conv2d/kernel', # , [1, 1, 32, 8]), 221 | model._blocks[i]._se_expand.bias: 'blocks_' + str(i) + '/se/conv2d_1/bias', # , [32]), 222 | model._blocks[i]._se_expand.weight: 'blocks_' + str(i) + '/se/conv2d_1/kernel', # , [1, 1, 8, 32]), 223 | model._blocks[i]._bn1.bias: 'blocks_' + str(i) + '/tpu_batch_normalization/beta', # [32]), 224 | model._blocks[i]._bn1.weight: 'blocks_' + str(i) + '/tpu_batch_normalization/gamma', # [32]), 225 | model._blocks[i]._bn1.running_mean: 'blocks_' + str(i) + '/tpu_batch_normalization/moving_mean', 226 | model._blocks[i]._bn1.running_var: 'blocks_' + str(i) + '/tpu_batch_normalization/moving_variance', 227 | model._blocks[i]._bn2.bias: 'blocks_' + str(i) + '/tpu_batch_normalization_1/beta', # [16]), 228 | model._blocks[i]._bn2.weight: 'blocks_' + str(i) + '/tpu_batch_normalization_1/gamma', # [16]), 229 | model._blocks[i]._bn2.running_mean: 'blocks_' + str(i) + '/tpu_batch_normalization_1/moving_mean', 230 | model._blocks[i]._bn2.running_var: 'blocks_' + str(i) + '/tpu_batch_normalization_1/moving_variance', 231 | } 232 | 233 | else: 234 | conversion_table_block = { 235 | model._blocks[i]._expand_conv.weight: 'blocks_' + str(i) + '/conv2d/kernel', 236 | model._blocks[i]._project_conv.weight: 'blocks_' + str(i) + '/conv2d_1/kernel', 237 | model._blocks[i]._depthwise_conv.weight: 'blocks_' + str(i) + '/depthwise_conv2d/depthwise_kernel', 238 | model._blocks[i]._se_reduce.bias: 'blocks_' + str(i) + '/se/conv2d/bias', 239 | model._blocks[i]._se_reduce.weight: 'blocks_' + str(i) + '/se/conv2d/kernel', 240 | model._blocks[i]._se_expand.bias: 'blocks_' + str(i) + '/se/conv2d_1/bias', 241 | model._blocks[i]._se_expand.weight: 'blocks_' + str(i) + '/se/conv2d_1/kernel', 242 | model._blocks[i]._bn0.bias: 'blocks_' + str(i) + '/tpu_batch_normalization/beta', 243 | model._blocks[i]._bn0.weight: 'blocks_' + str(i) + '/tpu_batch_normalization/gamma', 244 | model._blocks[i]._bn0.running_mean: 'blocks_' + str(i) + '/tpu_batch_normalization/moving_mean', 245 | model._blocks[i]._bn0.running_var: 'blocks_' + str(i) + '/tpu_batch_normalization/moving_variance', 246 | model._blocks[i]._bn1.bias: 'blocks_' + str(i) + '/tpu_batch_normalization_1/beta', 247 | model._blocks[i]._bn1.weight: 'blocks_' + str(i) + '/tpu_batch_normalization_1/gamma', 248 | model._blocks[i]._bn1.running_mean: 'blocks_' + str(i) + '/tpu_batch_normalization_1/moving_mean', 249 | model._blocks[i]._bn1.running_var: 'blocks_' + str(i) + '/tpu_batch_normalization_1/moving_variance', 250 | model._blocks[i]._bn2.bias: 'blocks_' + str(i) + '/tpu_batch_normalization_2/beta', 251 | model._blocks[i]._bn2.weight: 'blocks_' + str(i) + '/tpu_batch_normalization_2/gamma', 252 | model._blocks[i]._bn2.running_mean: 'blocks_' + str(i) + '/tpu_batch_normalization_2/moving_mean', 253 | model._blocks[i]._bn2.running_var: 'blocks_' + str(i) + '/tpu_batch_normalization_2/moving_variance', 254 | } 255 | 256 | conversion_table = merge(conversion_table, conversion_table_block) 257 | 258 | return conversion_table 259 | 260 | def load_and_save_temporary_tensorflow_model(model_name, model_ckpt, example_img= '../../example/img.jpg'): 261 | """ Loads and saves a TensorFlow model. """ 262 | image_files = [example_img] 263 | eval_ckpt_driver = eval_ckpt_main.EvalCkptDriver(model_name) 264 | with tf.Graph().as_default(), tf.Session() as sess: 265 | images, labels = eval_ckpt_driver.build_dataset(image_files, [0] * len(image_files), False) 266 | probs = eval_ckpt_driver.build_model(images, is_training=False) 267 | sess.run(tf.global_variables_initializer()) 268 | print(model_ckpt) 269 | eval_ckpt_driver.restore_model(sess, model_ckpt) 270 | tf.train.Saver().save(sess, 'tmp/model.ckpt') 271 | 272 | 273 | if __name__ == '__main__': 274 | 275 | import sys 276 | import argparse 277 | 278 | sys.path.append('original_tf') 279 | import eval_ckpt_main 280 | 281 | from model import * 282 | 283 | parser = argparse.ArgumentParser( 284 | description='Convert TF model to PyTorch model and save for easier future loading, \ 285 | python3 convert_params_tf_pytorch.py --model_name efficientnet-b0 \ 286 | --tf_checkpoint ../pretrained_tensorflow/efficientnet-b0/ \ 287 | --output_file ../pretrained_pytorch/efficientnet-b0.pth') 288 | 289 | parser.add_argument('--model_name', type=str, default='efficientnet-b0', 290 | help='efficientnet-b{N}, where N is an integer 0 <= N <= 7') 291 | parser.add_argument('--tf_checkpoint', type=str, default='../pretrained_tensorflow/efficientnet-b0/', 292 | help='checkpoint file path') 293 | parser.add_argument('--output_file', type=str, default='../pretrained_pytorch/efficientnet-b0.pth', 294 | help='output PyTorch model file name') 295 | args = parser.parse_args() 296 | 297 | # Build model 298 | model = get_from_name(args.model_name) 299 | 300 | # Load and save temporary TensorFlow file due to TF nuances 301 | print(args.tf_checkpoint) 302 | load_and_save_temporary_tensorflow_model(args.model_name, args.tf_checkpoint) 303 | 304 | pyt_tf_name_dict = pytorch_tf_name_dict(model) 305 | pyt_tf_params_dict = pytorch_tf_params_dict(model) 306 | pyt_params_dict = pytorch_param_from_tf_checkpoint('tmp/model.ckpt', pyt_tf_name_dict, pyt_tf_params_dict,args.model_name) 307 | 308 | # load pytorch param dict 309 | #pytorch_dict = pytorch_tf_dict(model,'tmp/model.ckpt',model_name=args.model_name) 310 | state_dict = model.state_dict() 311 | for k,v in state_dict.items(): 312 | if k in pyt_params_dict: 313 | print(k," : ",v.numpy().shape," --> ",pyt_params_dict[k].shape) 314 | state_dict[k] = torch.from_numpy(pyt_params_dict[k]) 315 | else: 316 | print("Warning --> ",k," not in pytorch_params_dict!!!") 317 | 318 | torch.save(state_dict, args.output_file) 319 | print('Saved model to', args.output_file) 320 | -------------------------------------------------------------------------------- /convert_tf_to_pt/model.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | from torch import nn 4 | from torch.nn import functional as F 5 | 6 | from utils import * 7 | 8 | class MBConvBlock(nn.Module): 9 | """ 10 | Mobile Inverted Residual Bottleneck Block 11 | 12 | Args: 13 | block_args (namedtuple): BlockArgs, see above 14 | global_params (namedtuple): GlobalParam, see above 15 | 16 | Attributes: 17 | has_se (bool): Whether the block contains a Squeeze and Excitation layer. 18 | """ 19 | 20 | def __init__(self, block_args, global_params,insize,drop_rate = None): 21 | super(MBConvBlock,self).__init__() 22 | print(block_args) 23 | self._block_args = block_args 24 | self._bn_mom = 1 - global_params.batch_norm_momentum 25 | self._bn_eps = global_params.batch_norm_epsilon 26 | self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1) 27 | self.id_skip = block_args.id_skip # skip connection and drop connect 28 | 29 | # Expansion phase 30 | outsize = insize 31 | inp = self._block_args.input_filters 32 | oup = self._block_args.input_filters * self._block_args.expand_ratio # number of output channels 33 | bsize = 6.0 34 | if self._block_args.expand_ratio != 1: 35 | self._expand_conv = getConv2d(in_size = insize, in_channels=inp,out_size = outsize, out_channels=oup,kernel_size = 1,bias = False) 36 | #self._expand_conv = Conv2dSamePadding(in_channels=inp, out_channels=oup, kernel_size=1, bias=False) 37 | self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) 38 | self._swish0 = Swish() 39 | bsize = 8.0 40 | 41 | # Depthwise convolution phase 42 | k = self._block_args.kernel_size 43 | s = self._block_args.stride 44 | outsize = insize/s[0] 45 | 46 | self._depthwise_conv = getConv2d(in_size = insize,in_channels=oup,out_size = outsize,out_channels=oup,groups=oup,kernel_size=k, stride=s[0], bias=False) 47 | #self._depthwise_conv = Conv2dSamePadding( 48 | # in_channels=oup, out_channels=oup, groups=oup, # groups makes it depthwise 49 | # kernel_size=k, stride=s, bias=False) 50 | self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) 51 | self._swish1 = Swish() 52 | # Squeeze and Excitation layer, if desired 53 | if self.has_se: 54 | num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio)) 55 | self._avg_pool = torch.nn.AdaptiveAvgPool2d(1) 56 | self._se_reduce = getConv2d(in_size = 1,in_channels=oup,out_size = 1,out_channels=num_squeezed_channels, kernel_size=1) 57 | self._se_expand = getConv2d(in_size = 1,in_channels=num_squeezed_channels,out_size = 1,out_channels=oup, kernel_size=1) 58 | self._swish2 = Swish() 59 | self._sigmoid = torch.nn.Sigmoid() 60 | self._mult = BroadcastMul() 61 | #self._se_reduce = Conv2dSamePadding(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1) 62 | #self._se_expand = Conv2dSamePadding(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1) 63 | 64 | # Output phase 65 | final_oup = self._block_args.output_filters 66 | self._project_conv = getConv2d(in_size = outsize,in_channels=oup,out_size = outsize,out_channels=final_oup, kernel_size=1, bias = False) 67 | #self._project_conv = Conv2dSamePadding(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False) 68 | self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps) 69 | 70 | input_filters, output_filters = self._block_args.input_filters, self._block_args.output_filters 71 | self.drop = drop_rate != None 72 | self.id_skip = self.id_skip and self._block_args.stride[0] == 1 and input_filters == output_filters 73 | print(self.drop,self.id_skip) 74 | if self.drop and self.id_skip: 75 | self.dropout = torch.nn.Dropout2d(drop_rate/bsize) 76 | if self.id_skip: 77 | self.shortcut = AddTensor() 78 | 79 | self.outsize = outsize 80 | 81 | def getOutSize(self): 82 | return self.outsize 83 | 84 | def forward(self, inputs): 85 | """ 86 | :param inputs: input tensor 87 | :param drop_connect_rate: drop connect rate (float, between 0 and 1) 88 | :return: output of block 89 | """ 90 | 91 | # Expansion and Depthwise Convolution 92 | x = inputs 93 | if self._block_args.expand_ratio != 1: 94 | x = self._swish0(self._bn0(self._expand_conv(inputs))) 95 | x = self._swish1(self._bn1(self._depthwise_conv(x))) 96 | 97 | # Squeeze and Excitation 98 | if self.has_se: 99 | x_squeezed = self._avg_pool(x)#F.adaptive_avg_pool2d(x, 1) 100 | x_squeezed = self._se_expand(self._swish2(self._se_reduce(x_squeezed))) 101 | x_squeezed = self._sigmoid(x_squeezed) 102 | x = self._mult(x,x_squeezed) 103 | 104 | x = self._bn2(self._project_conv(x)) 105 | if self.id_skip: 106 | if self.drop: 107 | x = self.dropout(x) 108 | x = self.shortcut(x,inputs) 109 | return x 110 | 111 | class EfficientNet(nn.Module): 112 | """ 113 | An EfficientNet model. Most easily loaded with the .from_name or .from_pretrained methods 114 | 115 | Args: 116 | blocks_args (list): A list of BlockArgs to construct blocks 117 | global_params (namedtuple): A set of GlobalParams shared between blocks 118 | 119 | Example: 120 | model = EfficientNet.from_pretrained('efficientnet-b0') 121 | 122 | """ 123 | def __init__(self, blocks_args=None, global_params=None): 124 | super(EfficientNet,self).__init__() 125 | assert isinstance(blocks_args, list), 'blocks_args should be a list' 126 | assert len(blocks_args) > 0, 'block args must be greater than 0' 127 | #print(blocks_args) 128 | self._global_params = global_params 129 | self._blocks_args = blocks_args 130 | 131 | # Batch norm parameters 132 | bn_mom = 1 - self._global_params.batch_norm_momentum 133 | bn_eps = self._global_params.batch_norm_epsilon 134 | 135 | # Stem 136 | insize = 224 137 | in_channels = 3 # rgb 138 | out_channels = round_filters(32, self._global_params) # number of output channels 139 | self._conv_stem = getConv2d(insize,in_channels,insize/2,out_channels,kernel_size=3, stride=2, bias=False)#Conv2dSamePadding(in_channels, out_channels, kernel_size=3, stride=2, bias=False) 140 | self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) 141 | self._swish0 = Swish() 142 | # Build blocks 143 | insize = insize/2 144 | self._blocks = nn.ModuleList([]) 145 | for block_args in self._blocks_args: 146 | #print(block_args) 147 | # Update block input and output filters based on depth multiplier. 148 | block_args = block_args._replace( 149 | input_filters=round_filters(block_args.input_filters, self._global_params), 150 | output_filters=round_filters(block_args.output_filters, self._global_params), 151 | num_repeat=round_repeats(block_args.num_repeat, self._global_params) 152 | ) 153 | 154 | # The first block needs to take care of stride and filter size increase. 155 | drop_rate = None 156 | if self._global_params.drop_connect_rate: 157 | drop_rate = self._global_params.drop_connect_rate*len(self._blocks) 158 | self._blocks.append(MBConvBlock(block_args, self._global_params,insize,drop_rate)) #### 159 | insize = self._blocks[-1].getOutSize() 160 | if block_args.num_repeat > 1: 161 | block_args = block_args._replace(input_filters=block_args.output_filters, stride=[1]) 162 | #print(block_args) 163 | for _ in range(block_args.num_repeat - 1): 164 | drop_rate = None 165 | if self._global_params.drop_connect_rate: 166 | drop_rate = self._global_params.drop_connect_rate * len(self._blocks) 167 | self._blocks.append(MBConvBlock(block_args, self._global_params,insize,drop_rate)) ### 168 | insize = self._blocks[-1].getOutSize() 169 | 170 | # Head 171 | in_channels = block_args.output_filters # output of final block 172 | out_channels = round_filters(1280, self._global_params) 173 | self._conv_head = getConv2d(insize,in_channels,insize,out_channels,kernel_size=1, bias=False)#Conv2dSamePadding(in_channels, out_channels, kernel_size=1, bias=False) 174 | self._bn1 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) 175 | self._swish1 = Swish() 176 | # Final linear layer 177 | self.global_pool = nn.AdaptiveAvgPool2d(1) 178 | self._dropout = self._global_params.dropout_rate 179 | if self._dropout: 180 | self.dropout = torch.nn.Dropout2d(self._dropout) 181 | self._fc = nn.Linear(out_channels, self._global_params.num_classes) 182 | 183 | def extract_features(self, inputs): 184 | """ Returns output of the final convolution layer """ 185 | # Stem 186 | x = self._swish0(self._bn0(self._conv_stem(inputs))) 187 | #x = self._conv_stem(inputs) 188 | #print x 189 | #x = self._conv_stem.weight 190 | #x = F.pad(x, [1,1,1,1]) 191 | #x = F.conv2d(inputs, self._conv_stem.weight, None, 2, 1, 1, 1) 192 | #x = inputs 193 | # Blocks 194 | for idx, block in enumerate(self._blocks): 195 | x = block(x) 196 | 197 | return x 198 | 199 | def forward(self, inputs): 200 | """ Calls extract_features to extract features, applies final linear layer, and returns logits. """ 201 | 202 | # Convolution layers 203 | x = self.extract_features(inputs) 204 | 205 | # Head 206 | x = self._swish1(self._bn1(self._conv_head(x))) 207 | x = self.global_pool(x).squeeze(-1).squeeze(-1) 208 | if self._dropout: 209 | x = self.dropout(x) 210 | x = F.dropout(x, p=self._dropout, training=self.training) 211 | x = self._fc(x) 212 | return x 213 | 214 | def get_from_name(model_name, override_params=None): 215 | #cls._check_model_name_is_valid(model_name) 216 | blocks_args, global_params = get_model_params(model_name, override_params) 217 | return EfficientNet(blocks_args, global_params) 218 | 219 | def get_from_pretrained(model_name): 220 | model = get_from_name(model_name) 221 | load_pretrained_weights(model, model_name) 222 | return model 223 | 224 | def get_image_size(cls, model_name): 225 | cls._check_model_name_is_valid(model_name) 226 | _, _, res, _ = efficientnet_params(model_name) 227 | return res 228 | 229 | def _check_model_name_is_valid(cls, model_name, also_need_pretrained_weights=False): 230 | """ Validates model name. None that pretrained weights are only available for 231 | the first four models (efficientnet-b{i} for i in 0,1,2,3) at the moment. """ 232 | num_models = 4 if also_need_pretrained_weights else 8 233 | valid_models = ['efficientnet_b'+str(i) for i in range(num_models)] 234 | if model_name.replace('-','_') not in valid_models: 235 | raise ValueError('model_name should be one of: ' + ', '.join(valid_models)) 236 | -------------------------------------------------------------------------------- /convert_tf_to_pt/model_caffe.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import caffe 4 | from torch import nn 5 | from torch.nn import functional as F 6 | 7 | from utils import * 8 | from prototxt_basic import * 9 | 10 | def caffe_active(ofs,top,bottom,type): 11 | attr = {} 12 | attr["top"] = top 13 | attr["bottom"] = [bottom] 14 | attr["attrs"] = {} 15 | attr["attrs"]["act_type"] = type 16 | Activation(ofs,attr) 17 | 18 | def caffe_normalization(ofs,top,bottom,eps): 19 | attr = {} 20 | attr["top"] = top 21 | attr["bottom"] = [bottom] 22 | attr["attrs"] = {} 23 | attr["attrs"]['eps'] = str(eps) 24 | BatchNorm(ofs,attr) 25 | 26 | def caffe_global_avg_pool(ofs,top,bottom): 27 | attr = {} 28 | attr["top"] = top 29 | attr["bottom"] = [bottom] 30 | attr["attrs"] = {} 31 | attr["attrs"]["pool_type"] = "avg" 32 | attr["attrs"]["global_pool"] = "True" 33 | Pooling_global(ofs,attr) 34 | 35 | def caffe_convolution(ofs,top,bottom,num_filter,kernel = 1,pad = 0,stride=1,bias=False,group = 1): 36 | attr = {} 37 | attr["top"] = top 38 | attr["bottom"] = [bottom] 39 | attr["attrs"] = {} 40 | attr["attrs"]["num_filter"] = str(num_filter) 41 | attr["attrs"]["kernel"] = "(" + str(kernel) + ",)" 42 | attr["attrs"]["pad"] = "(" + str(pad) + ",)" 43 | attr["attrs"]["stride"] = "(" + str(stride) + ",)" 44 | attr["attrs"]["num_group"] = str(group) 45 | attr["attrs"]["no_bias"] = str(not bias) 46 | Convolution(ofs,attr) 47 | 48 | def caffe_BroadcastMul(ofs,top,bottom): 49 | attr = {} 50 | attr["top"] = top 51 | attr["bottom"] = bottom 52 | broadcast_mul(ofs,attr) 53 | 54 | def caffe_dropout(ofs,top,bottom,drop_rate): 55 | attr = {} 56 | attr["top"] = top 57 | attr["bottom"] = [bottom] 58 | attr["attrs"] = {} 59 | attr["attrs"]["p"] = str(drop_rate) 60 | DropOut(ofs,attr) 61 | 62 | def caffe_elementwise(ofs,top,bottom): 63 | attr = {} 64 | attr["top"] = top 65 | attr["bottom"] = bottom 66 | ElementWiseSum(ofs,attr) 67 | 68 | def caffe_fc(ofs,top,bottom,num_hidden): 69 | attr = {} 70 | attr["top"] = top 71 | attr["bottom"] = [bottom] 72 | attr["attrs"] = {} 73 | attr["attrs"]["num_hidden"] = str(num_hidden) 74 | FullyConnected(ofs,attr) 75 | 76 | def caffe_flatten(ofs,top,bottom): 77 | attr = {} 78 | attr["top"] = top 79 | attr["bottom"] = [bottom] 80 | Flatten(ofs,attr) 81 | 82 | class MBConvBlock(nn.Module): 83 | """ 84 | Mobile Inverted Residual Bottleneck Block 85 | 86 | Args: 87 | block_args (namedtuple): BlockArgs, see above 88 | global_params (namedtuple): GlobalParam, see above 89 | 90 | Attributes: 91 | has_se (bool): Whether the block contains a Squeeze and Excitation layer. 92 | """ 93 | #caffe_attr {layerId,bottom_name,ofs} 94 | def __init__(self, block_args, global_params,caffe_attr,insize,drop_rate = None): 95 | super(MBConvBlock,self).__init__() 96 | print(block_args) 97 | #print(caffe_attr) 98 | self._block_args = block_args 99 | self._bn_mom = 1 - global_params.batch_norm_momentum 100 | self._bn_eps = global_params.batch_norm_epsilon 101 | self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1) 102 | self.id_skip = block_args.id_skip # skip connection and drop connect 103 | 104 | # Expansion phase 105 | outsize = insize 106 | inp = self._block_args.input_filters 107 | oup = self._block_args.input_filters * self._block_args.expand_ratio # number of output channels 108 | bsize = 6.0 109 | if self._block_args.expand_ratio != 1: 110 | caffe_convolution(caffe_attr[2],caffe_attr[0] + "._expand_conv",caffe_attr[1],oup) 111 | self._expand_conv = getConv2d(in_size = insize, in_channels=inp,out_size = outsize, out_channels=oup,kernel_size = 1,bias = False) 112 | #self._expand_conv = Conv2dSamePadding(in_channels=inp, out_channels=oup, kernel_size=1, bias=False) 113 | caffe_normalization(caffe_attr[2],caffe_attr[0] + "._bn0",caffe_attr[0] + "._expand_conv",self._bn_eps) 114 | self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) 115 | caffe_active(caffe_attr[2],caffe_attr[0] + "._swish0",caffe_attr[0] + "._bn0","swish") 116 | self._swish0 = Swish() 117 | bsize = 8.0 118 | 119 | # Depthwise convolution phase 120 | k = self._block_args.kernel_size 121 | s = self._block_args.stride 122 | outsize = insize/s[0] 123 | 124 | if self._block_args.expand_ratio == 1: 125 | bottom = caffe_attr[1] 126 | else: 127 | bottom = caffe_attr[0] + "._swish0" 128 | 129 | caffe_convolution(caffe_attr[2],caffe_attr[0] + "._depthwise_conv",bottom,oup,k,pad = k/2,stride = s[0],group = oup) 130 | self._depthwise_conv = getConv2d(in_size = insize,in_channels=oup,out_size = outsize,out_channels=oup,groups=oup,kernel_size=k, stride=s[0], bias=False) 131 | caffe_normalization(caffe_attr[2],caffe_attr[0] + "._bn1",caffe_attr[0] + "._depthwise_conv",self._bn_eps) 132 | self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) 133 | 134 | caffe_active(caffe_attr[2],caffe_attr[0] + "._swish1",caffe_attr[0] + "._bn1","swish") 135 | self._swish1 = Swish() 136 | # Squeeze and Excitation layer, if desired 137 | if self.has_se: 138 | num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio)) 139 | caffe_global_avg_pool(caffe_attr[2],caffe_attr[0] + "._global_avg_pool",caffe_attr[0] + "._swish1") 140 | self._global_avg_pool = torch.nn.AdaptiveAvgPool2d(1) 141 | caffe_convolution(caffe_attr[2],caffe_attr[0] + "._se_reduce",caffe_attr[0] + "._global_avg_pool",num_squeezed_channels,bias=True) 142 | self._se_reduce = getConv2d(in_size = 1,in_channels=oup,out_size = 1,out_channels=num_squeezed_channels, kernel_size=1) 143 | caffe_active(caffe_attr[2],caffe_attr[0] + "._swish2",caffe_attr[0] + "._se_reduce","swish") 144 | self._swish2 = Swish() 145 | caffe_convolution(caffe_attr[2],caffe_attr[0] + "._se_expand",caffe_attr[0] + "._swish2",oup,bias=True) 146 | self._se_expand = getConv2d(in_size = 1,in_channels=num_squeezed_channels,out_size = 1,out_channels=oup, kernel_size=1) 147 | caffe_active(caffe_attr[2],caffe_attr[0] + "._sigmoid",caffe_attr[0] + "._se_expand","sigmoid") 148 | self._sigmoid = torch.nn.Sigmoid() 149 | 150 | caffe_BroadcastMul(caffe_attr[2],caffe_attr[0] + "._broadcast_mul",[caffe_attr[0] + "._swish1",caffe_attr[0] + "._sigmoid"]) 151 | self._mult = BroadcastMul() 152 | #self._se_reduce = Conv2dSamePadding(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1) 153 | #self._se_expand = Conv2dSamePadding(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1) 154 | 155 | # Output phase 156 | final_oup = self._block_args.output_filters 157 | if self.has_se: 158 | bottom = caffe_attr[0] + "._broadcast_mul" 159 | else: 160 | bottom = caffe_attr[0] + "._swish1" 161 | 162 | caffe_convolution(caffe_attr[2],caffe_attr[0] + "._project_conv",bottom,final_oup) 163 | self._project_conv = getConv2d(in_size = outsize,in_channels=oup,out_size = outsize,out_channels=final_oup, kernel_size=1, bias = False) 164 | #self._project_conv = Conv2dSamePadding(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False) 165 | caffe_normalization(caffe_attr[2],caffe_attr[0] + "._bn2",caffe_attr[0] + "._project_conv",self._bn_eps) 166 | self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps) 167 | 168 | input_filters, output_filters = self._block_args.input_filters, self._block_args.output_filters 169 | self.drop = drop_rate != None 170 | self.id_skip = self.id_skip and self._block_args.stride[0] == 1 and input_filters == output_filters 171 | print(self.drop,self.id_skip) 172 | bottom = caffe_attr[0] + "._bn2" 173 | self.bottom_name = bottom 174 | if self.drop and self.id_skip: 175 | caffe_dropout(caffe_attr[2],caffe_attr[0] + "._dropout",caffe_attr[0] + "._bn2",drop_rate/bsize) 176 | self._dropout = torch.nn.Dropout2d(drop_rate/bsize) 177 | bottom = caffe_attr[0] + "._dropout" 178 | 179 | if self.id_skip: 180 | caffe_elementwise(caffe_attr[2],caffe_attr[0] + "._shortcut",[caffe_attr[1],bottom]) 181 | self._shortcut = AddTensor() 182 | self.bottom_name = caffe_attr[0] + "._shortcut" 183 | 184 | self.outsize = outsize 185 | 186 | def getOutSize(self): 187 | return self.outsize 188 | def getOutputName(self): 189 | return self.bottom_name 190 | 191 | def forward(self, inputs): 192 | """ 193 | :param inputs: input tensor 194 | :param drop_connect_rate: drop connect rate (float, between 0 and 1) 195 | :return: output of block 196 | """ 197 | 198 | # Expansion and Depthwise Convolution 199 | x = inputs 200 | if self._block_args.expand_ratio != 1: 201 | x = self._swish0(self._bn0(self._expand_conv(inputs))) 202 | x = self._swish1(self._bn1(self._depthwise_conv(x))) 203 | 204 | # Squeeze and Excitation 205 | if self.has_se: 206 | x_squeezed = self._global_avg_pool(x)#F.adaptive_avg_pool2d(x, 1) 207 | x_squeezed = self._se_expand(self._swish2(self._se_reduce(x_squeezed))) 208 | x_squeezed = self._sigmoid(x_squeezed) 209 | x = self._mult(x,x_squeezed) 210 | 211 | x = self._bn2(self._project_conv(x)) 212 | if self.id_skip: 213 | if self.drop: 214 | x = self._dropout(x) 215 | x = self._shortcut(x,inputs) 216 | return x 217 | 218 | class EfficientNet(nn.Module): 219 | """ 220 | An EfficientNet model. Most easily loaded with the .from_name or .from_pretrained methods 221 | 222 | Args: 223 | blocks_args (list): A list of BlockArgs to construct blocks 224 | global_params (namedtuple): A set of GlobalParams shared between blocks 225 | 226 | Example: 227 | model = EfficientNet.from_pretrained('efficientnet-b0') 228 | 229 | """ 230 | def __init__(self, blocks_args=None, global_params=None, model_name = "efficientnet-b0"): 231 | super(EfficientNet,self).__init__() 232 | assert isinstance(blocks_args, list), 'blocks_args should be a list' 233 | assert len(blocks_args) > 0, 'block args must be greater than 0' 234 | #print(blocks_args) 235 | self._global_params = global_params 236 | self._blocks_args = blocks_args 237 | 238 | # Batch norm parameters 239 | bn_mom = 1 - self._global_params.batch_norm_momentum 240 | bn_eps = self._global_params.batch_norm_epsilon 241 | 242 | # Stem 243 | insize = 224 244 | in_channels = 3 # rgb 245 | out_channels = round_filters(32, self._global_params) # number of output channels 246 | 247 | file = open("caffemodel/" + model_name + ".prototxt","w") 248 | data(file,{}) 249 | caffe_convolution(file,"_conv_stem","data",out_channels, 3, pad = 1 , stride = 2) 250 | self._conv_stem = getConv2d(insize,in_channels,insize/2,out_channels,kernel_size=3, stride=2, bias=False)#Conv2dSamePadding(in_channels, out_channels, kernel_size=3, stride=2, bias=False) 251 | caffe_normalization(file,"_bn0","_conv_stem",bn_eps) 252 | self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) 253 | caffe_active(file,"_swish0","_bn0","swish") 254 | self._swish0 = Swish() 255 | # Build blocks 256 | insize = insize/2 257 | self._blocks = nn.ModuleList([]) 258 | #print self._blocks_args 259 | last = "_swish0" 260 | for block_args in self._blocks_args: 261 | #print(block_args) 262 | # Update block input and output filters based on depth multiplier. 263 | block_args = block_args._replace( 264 | input_filters=round_filters(block_args.input_filters, self._global_params), 265 | output_filters=round_filters(block_args.output_filters, self._global_params), 266 | num_repeat=round_repeats(block_args.num_repeat, self._global_params) 267 | ) 268 | 269 | # The first block needs to take care of stride and filter size increase. 270 | drop_rate = None 271 | if self._global_params.drop_connect_rate: 272 | drop_rate = self._global_params.drop_connect_rate*len(self._blocks) 273 | self._blocks.append(MBConvBlock(block_args, self._global_params,("_blocks." + str(len(self._blocks)),last,file),insize,drop_rate)) #### 274 | insize = self._blocks[-1].getOutSize() 275 | last = self._blocks[-1].getOutputName() 276 | if block_args.num_repeat > 1: 277 | block_args = block_args._replace(input_filters=block_args.output_filters, stride=[1]) 278 | #print(block_args) 279 | for _ in range(block_args.num_repeat - 1): 280 | drop_rate = None 281 | if self._global_params.drop_connect_rate: 282 | drop_rate = self._global_params.drop_connect_rate * len(self._blocks) 283 | self._blocks.append(MBConvBlock(block_args, self._global_params,("_blocks." + str(len(self._blocks)),last,file),insize,drop_rate)) ### 284 | insize = self._blocks[-1].getOutSize() 285 | last = self._blocks[-1].getOutputName() 286 | # Head 287 | in_channels = block_args.output_filters # output of final block 288 | out_channels = round_filters(1280, self._global_params) 289 | caffe_convolution(file,"_conv_head",last,out_channels) 290 | self._conv_head = getConv2d(insize,in_channels,insize,out_channels,kernel_size=1, bias=False)#Conv2dSamePadding(in_channels, out_channels, kernel_size=1, bias=False) 291 | caffe_normalization(file,"_bn1","_conv_head",bn_eps) 292 | self._bn1 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) 293 | caffe_active(file,"_swish1","_bn1","swish") 294 | self._swish1 = Swish() 295 | # Final linear layer 296 | caffe_global_avg_pool(file,"_global_avg_pool","_swish1") 297 | self._global_avg_pool = nn.AdaptiveAvgPool2d(1) 298 | self._dropout = self._global_params.dropout_rate 299 | if self._dropout: 300 | caffe_dropout(file,"_dropout","_global_avg_pool",self._dropout) 301 | self._dropout = torch.nn.Dropout2d(self._dropout) 302 | caffe_flatten(file,"_flatten","_dropout") 303 | caffe_fc(file,"_fc","_flatten",self._global_params.num_classes) 304 | self._fc = nn.Linear(out_channels, self._global_params.num_classes) 305 | file.close() 306 | 307 | def extract_features(self, inputs): 308 | """ Returns output of the final convolution layer """ 309 | # Stem 310 | x = self._swish0(self._bn0(self._conv_stem(inputs))) 311 | 312 | # Blocks 313 | for idx, block in enumerate(self._blocks): 314 | x = block(x) 315 | return x 316 | 317 | def forward(self, inputs): 318 | """ Calls extract_features to extract features, applies final linear layer, and returns logits. """ 319 | 320 | # Convolution layers 321 | x = self.extract_features(inputs) 322 | 323 | # Head 324 | x = self._swish1(self._bn1(self._conv_head(x))) 325 | #x = F.adaptive_avg_pool2d(x, 1).squeeze(-1).squeeze(-1) 326 | x = self._global_avg_pool(x).squeeze(-1).squeeze(-1) 327 | if self._dropout: 328 | x = self.dropout(x) 329 | #x = F.dropout(x, p=self._dropout, training=self.training) 330 | x = self._fc(x) 331 | return x 332 | 333 | def get_from_name(model_name, override_params=None): 334 | #cls._check_model_name_is_valid(model_name) 335 | blocks_args, global_params = get_model_params(model_name, override_params) 336 | return EfficientNet(blocks_args, global_params,model_name) 337 | 338 | def get_from_pretrained(model_name): 339 | model = EfficientNet.from_name(model_name) 340 | load_pretrained_weights(model, model_name) 341 | return model 342 | 343 | def get_image_size(cls, model_name): 344 | cls._check_model_name_is_valid(model_name) 345 | _, _, res, _ = efficientnet_params(model_name) 346 | return res 347 | 348 | def _check_model_name_is_valid(cls, model_name, also_need_pretrained_weights=False): 349 | """ Validates model name. None that pretrained weights are only available for 350 | the first four models (efficientnet-b{i} for i in 0,1,2,3) at the moment. """ 351 | num_models = 4 if also_need_pretrained_weights else 8 352 | valid_models = ['efficientnet_b'+str(i) for i in range(num_models)] 353 | if model_name.replace('-','_') not in valid_models: 354 | raise ValueError('model_name should be one of: ' + ', '.join(valid_models)) 355 | -------------------------------------------------------------------------------- /convert_tf_to_pt/original_tf/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxworkspace/efficientnet_tf_pytorch_caffe/3e683de7e7667165e0c1cb60d3b0ca68c8ba9098/convert_tf_to_pt/original_tf/__init__.py -------------------------------------------------------------------------------- /convert_tf_to_pt/original_tf/efficientnet_builder.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Model Builder for EfficientNet.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import os 22 | import re 23 | import tensorflow as tf 24 | 25 | import efficientnet_model 26 | 27 | 28 | def efficientnet_params(model_name): 29 | """Get efficientnet params based on model name.""" 30 | params_dict = { 31 | # (width_coefficient, depth_coefficient, resolution, dropout_rate) 32 | 'efficientnet-b0': (1.0, 1.0, 224, 0.2), 33 | 'efficientnet-b1': (1.0, 1.1, 240, 0.2), 34 | 'efficientnet-b2': (1.1, 1.2, 260, 0.3), 35 | 'efficientnet-b3': (1.2, 1.4, 300, 0.3), 36 | 'efficientnet-b4': (1.4, 1.8, 380, 0.4), 37 | 'efficientnet-b5': (1.6, 2.2, 456, 0.4), 38 | 'efficientnet-b6': (1.8, 2.6, 528, 0.5), 39 | 'efficientnet-b7': (2.0, 3.1, 600, 0.5), 40 | } 41 | return params_dict[model_name] 42 | 43 | 44 | class BlockDecoder(object): 45 | """Block Decoder for readability.""" 46 | 47 | def _decode_block_string(self, block_string): 48 | """Gets a block through a string notation of arguments.""" 49 | assert isinstance(block_string, str) 50 | ops = block_string.split('_') 51 | options = {} 52 | for op in ops: 53 | splits = re.split(r'(\d.*)', op) 54 | if len(splits) >= 2: 55 | key, value = splits[:2] 56 | options[key] = value 57 | 58 | if 's' not in options or len(options['s']) != 2: 59 | raise ValueError('Strides options should be a pair of integers.') 60 | 61 | return efficientnet_model.BlockArgs( 62 | kernel_size=int(options['k']), 63 | num_repeat=int(options['r']), 64 | input_filters=int(options['i']), 65 | output_filters=int(options['o']), 66 | expand_ratio=int(options['e']), 67 | id_skip=('noskip' not in block_string), 68 | se_ratio=float(options['se']) if 'se' in options else None, 69 | strides=[int(options['s'][0]), int(options['s'][1])]) 70 | 71 | def _encode_block_string(self, block): 72 | """Encodes a block to a string.""" 73 | args = [ 74 | 'r%d' % block.num_repeat, 75 | 'k%d' % block.kernel_size, 76 | 's%d%d' % (block.strides[0], block.strides[1]), 77 | 'e%s' % block.expand_ratio, 78 | 'i%d' % block.input_filters, 79 | 'o%d' % block.output_filters 80 | ] 81 | if block.se_ratio > 0 and block.se_ratio <= 1: 82 | args.append('se%s' % block.se_ratio) 83 | if block.id_skip is False: 84 | args.append('noskip') 85 | return '_'.join(args) 86 | 87 | def decode(self, string_list): 88 | """Decodes a list of string notations to specify blocks inside the network. 89 | 90 | Args: 91 | string_list: a list of strings, each string is a notation of block. 92 | 93 | Returns: 94 | A list of namedtuples to represent blocks arguments. 95 | """ 96 | assert isinstance(string_list, list) 97 | blocks_args = [] 98 | for block_string in string_list: 99 | blocks_args.append(self._decode_block_string(block_string)) 100 | return blocks_args 101 | 102 | def encode(self, blocks_args): 103 | """Encodes a list of Blocks to a list of strings. 104 | 105 | Args: 106 | blocks_args: A list of namedtuples to represent blocks arguments. 107 | Returns: 108 | a list of strings, each string is a notation of block. 109 | """ 110 | block_strings = [] 111 | for block in blocks_args: 112 | block_strings.append(self._encode_block_string(block)) 113 | return block_strings 114 | 115 | 116 | def efficientnet(width_coefficient=None, 117 | depth_coefficient=None, 118 | dropout_rate=0.2, 119 | drop_connect_rate=0.2): 120 | """Creates a efficientnet model.""" 121 | blocks_args = [ 122 | 'r1_k3_s11_e1_i32_o16_se0.25', 'r2_k3_s22_e6_i16_o24_se0.25', 123 | 'r2_k5_s22_e6_i24_o40_se0.25', 'r3_k3_s22_e6_i40_o80_se0.25', 124 | 'r3_k5_s11_e6_i80_o112_se0.25', 'r4_k5_s22_e6_i112_o192_se0.25', 125 | 'r1_k3_s11_e6_i192_o320_se0.25', 126 | ] 127 | global_params = efficientnet_model.GlobalParams( 128 | batch_norm_momentum=0.99, 129 | batch_norm_epsilon=1e-3, 130 | dropout_rate=dropout_rate, 131 | drop_connect_rate=drop_connect_rate, 132 | data_format='channels_last', 133 | num_classes=1000, 134 | width_coefficient=width_coefficient, 135 | depth_coefficient=depth_coefficient, 136 | depth_divisor=8, 137 | min_depth=None) 138 | decoder = BlockDecoder() 139 | return decoder.decode(blocks_args), global_params 140 | 141 | 142 | def get_model_params(model_name, override_params): 143 | """Get the block args and global params for a given model.""" 144 | if model_name.startswith('efficientnet'): 145 | width_coefficient, depth_coefficient, _, dropout_rate = ( 146 | efficientnet_params(model_name)) 147 | blocks_args, global_params = efficientnet( 148 | width_coefficient, depth_coefficient, dropout_rate) 149 | else: 150 | raise NotImplementedError('model name is not pre-defined: %s' % model_name) 151 | 152 | if override_params: 153 | # ValueError will be raised here if override_params has fields not included 154 | # in global_params. 155 | global_params = global_params._replace(**override_params) 156 | 157 | tf.logging.info('global_params= %s', global_params) 158 | tf.logging.info('blocks_args= %s', blocks_args) 159 | return blocks_args, global_params 160 | 161 | 162 | def build_model(images, 163 | model_name, 164 | training, 165 | override_params=None, 166 | model_dir=None): 167 | """A helper functiion to creates a model and returns predicted logits. 168 | 169 | Args: 170 | images: input images tensor. 171 | model_name: string, the predefined model name. 172 | training: boolean, whether the model is constructed for training. 173 | override_params: A dictionary of params for overriding. Fields must exist in 174 | efficientnet_model.GlobalParams. 175 | model_dir: string, optional model dir for saving configs. 176 | 177 | Returns: 178 | logits: the logits tensor of classes. 179 | endpoints: the endpoints for each layer. 180 | 181 | Raises: 182 | When model_name specified an undefined model, raises NotImplementedError. 183 | When override_params has invalid fields, raises ValueError. 184 | """ 185 | assert isinstance(images, tf.Tensor) 186 | blocks_args, global_params = get_model_params(model_name, override_params) 187 | 188 | if model_dir: 189 | param_file = os.path.join(model_dir, 'model_params.txt') 190 | if not tf.gfile.Exists(param_file): 191 | with tf.gfile.GFile(param_file, 'w') as f: 192 | tf.logging.info('writing to %s' % param_file) 193 | f.write('model_name= %s\n\n' % model_name) 194 | f.write('global_params= %s\n\n' % str(global_params)) 195 | f.write('blocks_args= %s\n\n' % str(blocks_args)) 196 | 197 | with tf.variable_scope(model_name): 198 | model = efficientnet_model.Model(blocks_args, global_params) 199 | logits = model(images, training=training) 200 | 201 | logits = tf.identity(logits, 'logits') 202 | return logits, model.endpoints 203 | 204 | 205 | def build_model_base(images, model_name, training, override_params=None): 206 | """A helper functiion to create a base model and return global_pool. 207 | 208 | Args: 209 | images: input images tensor. 210 | model_name: string, the model name of a pre-defined MnasNet. 211 | training: boolean, whether the model is constructed for training. 212 | override_params: A dictionary of params for overriding. Fields must exist in 213 | mnasnet_model.GlobalParams. 214 | 215 | Returns: 216 | features: global pool features. 217 | endpoints: the endpoints for each layer. 218 | 219 | Raises: 220 | When model_name specified an undefined model, raises NotImplementedError. 221 | When override_params has invalid fields, raises ValueError. 222 | """ 223 | assert isinstance(images, tf.Tensor) 224 | blocks_args, global_params = get_model_params(model_name, override_params) 225 | 226 | with tf.variable_scope(model_name): 227 | model = efficientnet_model.Model(blocks_args, global_params) 228 | features = model(images, training=training, features_only=True) 229 | 230 | features = tf.identity(features, 'global_pool') 231 | return features, model.endpoints 232 | -------------------------------------------------------------------------------- /convert_tf_to_pt/original_tf/efficientnet_model.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Contains definitions for EfficientNet model. 16 | 17 | [1] Mingxing Tan, Quoc V. Le 18 | EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks. 19 | ICML'19, https://arxiv.org/abs/1905.11946 20 | """ 21 | 22 | from __future__ import absolute_import 23 | from __future__ import division 24 | from __future__ import print_function 25 | 26 | import collections 27 | import math 28 | import numpy as np 29 | import six 30 | from six.moves import xrange # pylint: disable=redefined-builtin 31 | import tensorflow as tf 32 | 33 | #from efficientnet_pytorch import utils 34 | from original_tf import utils 35 | 36 | GlobalParams = collections.namedtuple('GlobalParams', [ 37 | 'batch_norm_momentum', 'batch_norm_epsilon', 'dropout_rate', 'data_format', 38 | 'num_classes', 'width_coefficient', 'depth_coefficient', 39 | 'depth_divisor', 'min_depth', 'drop_connect_rate', 40 | ]) 41 | GlobalParams.__new__.__defaults__ = (None,) * len(GlobalParams._fields) 42 | 43 | # batchnorm = tf.layers.BatchNormalization 44 | batchnorm = utils.TpuBatchNormalization # TPU-specific requirement. 45 | relu_fn = tf.nn.swish 46 | 47 | 48 | BlockArgs = collections.namedtuple('BlockArgs', [ 49 | 'kernel_size', 'num_repeat', 'input_filters', 'output_filters', 50 | 'expand_ratio', 'id_skip', 'strides', 'se_ratio' 51 | ]) 52 | # defaults will be a public argument for namedtuple in Python 3.7 53 | # https://docs.python.org/3/library/collections.html#collections.namedtuple 54 | BlockArgs.__new__.__defaults__ = (None,) * len(BlockArgs._fields) 55 | 56 | 57 | def conv_kernel_initializer(shape, dtype=None, partition_info=None): 58 | """Initialization for convolutional kernels. 59 | 60 | The main difference with tf.variance_scaling_initializer is that 61 | tf.variance_scaling_initializer uses a truncated normal with an uncorrected 62 | standard deviation, whereas here we use a normal distribution. Similarly, 63 | tf.contrib.layers.variance_scaling_initializer uses a truncated normal with 64 | a corrected standard deviation. 65 | 66 | Args: 67 | shape: shape of variable 68 | dtype: dtype of variable 69 | partition_info: unused 70 | 71 | Returns: 72 | an initialization for the variable 73 | """ 74 | del partition_info 75 | kernel_height, kernel_width, _, out_filters = shape 76 | fan_out = int(kernel_height * kernel_width * out_filters) 77 | return tf.random_normal( 78 | shape, mean=0.0, stddev=np.sqrt(2.0 / fan_out), dtype=dtype) 79 | 80 | 81 | def dense_kernel_initializer(shape, dtype=None, partition_info=None): 82 | """Initialization for dense kernels. 83 | 84 | This initialization is equal to 85 | tf.variance_scaling_initializer(scale=1.0/3.0, mode='fan_out', 86 | distribution='uniform'). 87 | It is written out explicitly here for clarity. 88 | 89 | Args: 90 | shape: shape of variable 91 | dtype: dtype of variable 92 | partition_info: unused 93 | 94 | Returns: 95 | an initialization for the variable 96 | """ 97 | del partition_info 98 | init_range = 1.0 / np.sqrt(shape[1]) 99 | return tf.random_uniform(shape, -init_range, init_range, dtype=dtype) 100 | 101 | 102 | def round_filters(filters, global_params): 103 | """Round number of filters based on depth multiplier.""" 104 | orig_f = filters 105 | multiplier = global_params.width_coefficient 106 | divisor = global_params.depth_divisor 107 | min_depth = global_params.min_depth 108 | if not multiplier: 109 | return filters 110 | 111 | filters *= multiplier 112 | min_depth = min_depth or divisor 113 | new_filters = max(min_depth, int(filters + divisor / 2) // divisor * divisor) 114 | # Make sure that round down does not go down by more than 10%. 115 | if new_filters < 0.9 * filters: 116 | new_filters += divisor 117 | tf.logging.info('round_filter input={} output={}'.format(orig_f, new_filters)) 118 | return int(new_filters) 119 | 120 | 121 | def round_repeats(repeats, global_params): 122 | """Round number of filters based on depth multiplier.""" 123 | multiplier = global_params.depth_coefficient 124 | if not multiplier: 125 | return repeats 126 | return int(math.ceil(multiplier * repeats)) 127 | 128 | 129 | class MBConvBlock(object): 130 | """A class of MBConv: Mobile Inveretd Residual Bottleneck. 131 | 132 | Attributes: 133 | has_se: boolean. Whether the block contains a Squeeze and Excitation layer 134 | inside. 135 | endpoints: dict. A list of internal tensors. 136 | """ 137 | 138 | def __init__(self, block_args, global_params): 139 | """Initializes a MBConv block. 140 | 141 | Args: 142 | block_args: BlockArgs, arguments to create a Block. 143 | global_params: GlobalParams, a set of global parameters. 144 | """ 145 | self._block_args = block_args 146 | self._batch_norm_momentum = global_params.batch_norm_momentum 147 | self._batch_norm_epsilon = global_params.batch_norm_epsilon 148 | if global_params.data_format == 'channels_first': 149 | self._channel_axis = 1 150 | self._spatial_dims = [2, 3] 151 | else: 152 | self._channel_axis = -1 153 | self._spatial_dims = [1, 2] 154 | self.has_se = (self._block_args.se_ratio is not None) and ( 155 | self._block_args.se_ratio > 0) and (self._block_args.se_ratio <= 1) 156 | 157 | self.endpoints = None 158 | 159 | # Builds the block accordings to arguments. 160 | self._build() 161 | 162 | def block_args(self): 163 | return self._block_args 164 | 165 | def _build(self): 166 | """Builds block according to the arguments.""" 167 | filters = self._block_args.input_filters * self._block_args.expand_ratio 168 | if self._block_args.expand_ratio != 1: 169 | # Expansion phase: 170 | self._expand_conv = tf.layers.Conv2D( 171 | filters, 172 | kernel_size=[1, 1], 173 | strides=[1, 1], 174 | kernel_initializer=conv_kernel_initializer, 175 | padding='same', 176 | use_bias=False) 177 | self._bn0 = batchnorm( 178 | axis=self._channel_axis, 179 | momentum=self._batch_norm_momentum, 180 | epsilon=self._batch_norm_epsilon) 181 | 182 | kernel_size = self._block_args.kernel_size 183 | # Depth-wise convolution phase: 184 | self._depthwise_conv = utils.DepthwiseConv2D( 185 | [kernel_size, kernel_size], 186 | strides=self._block_args.strides, 187 | depthwise_initializer=conv_kernel_initializer, 188 | padding='same', 189 | use_bias=False) 190 | self._bn1 = batchnorm( 191 | axis=self._channel_axis, 192 | momentum=self._batch_norm_momentum, 193 | epsilon=self._batch_norm_epsilon) 194 | 195 | if self.has_se: 196 | num_reduced_filters = max( 197 | 1, int(self._block_args.input_filters * self._block_args.se_ratio)) 198 | # Squeeze and Excitation layer. 199 | self._se_reduce = tf.layers.Conv2D( 200 | num_reduced_filters, 201 | kernel_size=[1, 1], 202 | strides=[1, 1], 203 | kernel_initializer=conv_kernel_initializer, 204 | padding='same', 205 | use_bias=True) 206 | self._se_expand = tf.layers.Conv2D( 207 | filters, 208 | kernel_size=[1, 1], 209 | strides=[1, 1], 210 | kernel_initializer=conv_kernel_initializer, 211 | padding='same', 212 | use_bias=True) 213 | 214 | # Output phase: 215 | filters = self._block_args.output_filters 216 | self._project_conv = tf.layers.Conv2D( 217 | filters, 218 | kernel_size=[1, 1], 219 | strides=[1, 1], 220 | kernel_initializer=conv_kernel_initializer, 221 | padding='same', 222 | use_bias=False) 223 | self._bn2 = batchnorm( 224 | axis=self._channel_axis, 225 | momentum=self._batch_norm_momentum, 226 | epsilon=self._batch_norm_epsilon) 227 | 228 | def _call_se(self, input_tensor): 229 | """Call Squeeze and Excitation layer. 230 | 231 | Args: 232 | input_tensor: Tensor, a single input tensor for Squeeze/Excitation layer. 233 | 234 | Returns: 235 | A output tensor, which should have the same shape as input. 236 | """ 237 | se_tensor = tf.reduce_mean(input_tensor, self._spatial_dims, keepdims=True) 238 | se_tensor = self._se_expand(relu_fn(self._se_reduce(se_tensor))) 239 | tf.logging.info('Built Squeeze and Excitation with tensor shape: %s' % 240 | (se_tensor.shape)) 241 | return tf.sigmoid(se_tensor) * input_tensor 242 | 243 | def call(self, inputs, training=True, drop_connect_rate=None): 244 | """Implementation of call(). 245 | 246 | Args: 247 | inputs: the inputs tensor. 248 | training: boolean, whether the model is constructed for training. 249 | drop_connect_rate: float, between 0 to 1, drop connect rate. 250 | 251 | Returns: 252 | A output tensor. 253 | """ 254 | tf.logging.info('Block input: %s shape: %s' % (inputs.name, inputs.shape)) 255 | if self._block_args.expand_ratio != 1: 256 | x = relu_fn(self._bn0(self._expand_conv(inputs), training=training)) 257 | else: 258 | x = inputs 259 | tf.logging.info('Expand: %s shape: %s' % (x.name, x.shape)) 260 | 261 | x = relu_fn(self._bn1(self._depthwise_conv(x), training=training)) 262 | tf.logging.info('DWConv: %s shape: %s' % (x.name, x.shape)) 263 | 264 | if self.has_se: 265 | with tf.variable_scope('se'): 266 | x = self._call_se(x) 267 | 268 | self.endpoints = {'expansion_output': x} 269 | 270 | x = self._bn2(self._project_conv(x), training=training) 271 | if self._block_args.id_skip: 272 | if all( 273 | s == 1 for s in self._block_args.strides 274 | ) and self._block_args.input_filters == self._block_args.output_filters: 275 | # only apply drop_connect if skip presents. 276 | if drop_connect_rate: 277 | x = utils.drop_connect(x, training, drop_connect_rate) 278 | x = tf.add(x, inputs) 279 | tf.logging.info('Project: %s shape: %s' % (x.name, x.shape)) 280 | return x 281 | 282 | 283 | class Model(tf.keras.Model): 284 | """A class implements tf.keras.Model for MNAS-like model. 285 | 286 | Reference: https://arxiv.org/abs/1807.11626 287 | """ 288 | 289 | def __init__(self, blocks_args=None, global_params=None): 290 | """Initializes an `Model` instance. 291 | 292 | Args: 293 | blocks_args: A list of BlockArgs to construct block modules. 294 | global_params: GlobalParams, a set of global parameters. 295 | 296 | Raises: 297 | ValueError: when blocks_args is not specified as a list. 298 | """ 299 | super(Model, self).__init__() 300 | if not isinstance(blocks_args, list): 301 | raise ValueError('blocks_args should be a list.') 302 | self._global_params = global_params 303 | self._blocks_args = blocks_args 304 | self.endpoints = None 305 | self._build() 306 | 307 | def _build(self): 308 | """Builds a model.""" 309 | self._blocks = [] 310 | # Builds blocks. 311 | for block_args in self._blocks_args: 312 | assert block_args.num_repeat > 0 313 | # Update block input and output filters based on depth multiplier. 314 | block_args = block_args._replace( 315 | input_filters=round_filters(block_args.input_filters, 316 | self._global_params), 317 | output_filters=round_filters(block_args.output_filters, 318 | self._global_params), 319 | num_repeat=round_repeats(block_args.num_repeat, self._global_params)) 320 | 321 | # The first block needs to take care of stride and filter size increase. 322 | self._blocks.append(MBConvBlock(block_args, self._global_params)) 323 | if block_args.num_repeat > 1: 324 | # pylint: disable=protected-access 325 | block_args = block_args._replace( 326 | input_filters=block_args.output_filters, strides=[1, 1]) 327 | # pylint: enable=protected-access 328 | for _ in xrange(block_args.num_repeat - 1): 329 | self._blocks.append(MBConvBlock(block_args, self._global_params)) 330 | 331 | batch_norm_momentum = self._global_params.batch_norm_momentum 332 | batch_norm_epsilon = self._global_params.batch_norm_epsilon 333 | if self._global_params.data_format == 'channels_first': 334 | channel_axis = 1 335 | else: 336 | channel_axis = -1 337 | 338 | # Stem part. 339 | self._conv_stem = tf.layers.Conv2D( 340 | filters=round_filters(32, self._global_params), 341 | kernel_size=[3, 3], 342 | strides=[2, 2], 343 | kernel_initializer=conv_kernel_initializer, 344 | padding='same', 345 | use_bias=False) 346 | self._bn0 = batchnorm( 347 | axis=channel_axis, 348 | momentum=batch_norm_momentum, 349 | epsilon=batch_norm_epsilon) 350 | 351 | # Head part. 352 | self._conv_head = tf.layers.Conv2D( 353 | filters=round_filters(1280, self._global_params), 354 | kernel_size=[1, 1], 355 | strides=[1, 1], 356 | kernel_initializer=conv_kernel_initializer, 357 | padding='same', 358 | use_bias=False) 359 | self._bn1 = batchnorm( 360 | axis=channel_axis, 361 | momentum=batch_norm_momentum, 362 | epsilon=batch_norm_epsilon) 363 | 364 | self._avg_pooling = tf.keras.layers.GlobalAveragePooling2D( 365 | data_format=self._global_params.data_format) 366 | self._fc = tf.layers.Dense( 367 | self._global_params.num_classes, 368 | kernel_initializer=dense_kernel_initializer) 369 | 370 | if self._global_params.dropout_rate > 0: 371 | self._dropout = tf.keras.layers.Dropout(self._global_params.dropout_rate) 372 | else: 373 | self._dropout = None 374 | 375 | def call(self, inputs, training=True, features_only=None): 376 | """Implementation of call(). 377 | 378 | Args: 379 | inputs: input tensors. 380 | training: boolean, whether the model is constructed for training. 381 | features_only: build the base feature network only. 382 | 383 | Returns: 384 | output tensors. 385 | """ 386 | outputs = None 387 | self.endpoints = {} 388 | # Calls Stem layers 389 | with tf.variable_scope('stem'): 390 | outputs = relu_fn( 391 | self._bn0(self._conv_stem(inputs), training=training)) 392 | tf.logging.info('Built stem layers with output shape: %s' % outputs.shape) 393 | self.endpoints['stem'] = outputs 394 | 395 | # Calls blocks. 396 | reduction_idx = 0 397 | for idx, block in enumerate(self._blocks): 398 | is_reduction = False 399 | if ((idx == len(self._blocks) - 1) or 400 | self._blocks[idx + 1].block_args().strides[0] > 1): 401 | is_reduction = True 402 | reduction_idx += 1 403 | 404 | with tf.variable_scope('blocks_%s' % idx): 405 | drop_rate = self._global_params.drop_connect_rate 406 | if drop_rate: 407 | drop_rate *= float(idx) / len(self._blocks) 408 | tf.logging.info('block_%s drop_connect_rate: %s' % (idx, drop_rate)) 409 | outputs = block.call(outputs, training=training) 410 | self.endpoints['block_%s' % idx] = outputs 411 | if is_reduction: 412 | self.endpoints['reduction_%s' % reduction_idx] = outputs 413 | if block.endpoints: 414 | for k, v in six.iteritems(block.endpoints): 415 | self.endpoints['block_%s/%s' % (idx, k)] = v 416 | if is_reduction: 417 | self.endpoints['reduction_%s/%s' % (reduction_idx, k)] = v 418 | self.endpoints['global_pool'] = outputs 419 | 420 | if not features_only: 421 | # Calls final layers and returns logits. 422 | with tf.variable_scope('head'): 423 | outputs = relu_fn( 424 | self._bn1(self._conv_head(outputs), training=training)) 425 | outputs = self._avg_pooling(outputs) 426 | if self._dropout: 427 | outputs = self._dropout(outputs, training=training) 428 | outputs = self._fc(outputs) 429 | self.endpoints['head'] = outputs 430 | return outputs 431 | -------------------------------------------------------------------------------- /convert_tf_to_pt/original_tf/eval_ckpt_main.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Eval checkpoint driver. 16 | 17 | This is an example evaluation script for users to understand the EfficientNet 18 | model checkpoints on CPU. To serve EfficientNet, please consider to export a 19 | `SavedModel` from checkpoints and use tf-serving to serve. 20 | """ 21 | 22 | from __future__ import absolute_import 23 | from __future__ import division 24 | from __future__ import print_function 25 | 26 | import json 27 | import sys 28 | from absl import app 29 | from absl import flags 30 | import numpy as np 31 | import tensorflow as tf 32 | 33 | 34 | import efficientnet_builder 35 | import preprocessing 36 | 37 | 38 | flags.DEFINE_string('model_name', 'efficientnet-b0', 'Model name to eval.') 39 | flags.DEFINE_string('runmode', 'examples', 'Running mode: examples or imagenet') 40 | flags.DEFINE_string('imagenet_eval_glob', None, 41 | 'Imagenet eval image glob, ' 42 | 'such as /imagenet/ILSVRC2012*.JPEG') 43 | flags.DEFINE_string('imagenet_eval_label', None, 44 | 'Imagenet eval label file path, ' 45 | 'such as /imagenet/ILSVRC2012_validation_ground_truth.txt') 46 | flags.DEFINE_string('ckpt_dir', '/tmp/ckpt/', 'Checkpoint folders') 47 | flags.DEFINE_string('example_img', '/tmp/panda.jpg', 48 | 'Filepath for a single example image.') 49 | flags.DEFINE_string('labels_map_file', '/tmp/labels_map.txt', 50 | 'Labels map from label id to its meaning.') 51 | flags.DEFINE_integer('num_images', 5000, 52 | 'Number of images to eval. Use -1 to eval all images.') 53 | FLAGS = flags.FLAGS 54 | 55 | MEAN_RGB = [0.485 * 255, 0.456 * 255, 0.406 * 255] 56 | STDDEV_RGB = [0.229 * 255, 0.224 * 255, 0.225 * 255] 57 | 58 | 59 | class EvalCkptDriver(object): 60 | """A driver for running eval inference. 61 | 62 | Attributes: 63 | model_name: str. Model name to eval. 64 | batch_size: int. Eval batch size. 65 | num_classes: int. Number of classes, default to 1000 for ImageNet. 66 | image_size: int. Input image size, determined by model name. 67 | """ 68 | 69 | def __init__(self, model_name='efficientnet-b0', batch_size=1): 70 | """Initialize internal variables.""" 71 | self.model_name = model_name 72 | self.batch_size = batch_size 73 | self.num_classes = 1000 74 | # Model Scaling parameters 75 | _, _, self.image_size, _ = efficientnet_builder.efficientnet_params( 76 | model_name) 77 | 78 | def restore_model(self, sess, ckpt_dir): 79 | """Restore variables from checkpoint dir.""" 80 | checkpoint = tf.train.latest_checkpoint(ckpt_dir) 81 | ema = tf.train.ExponentialMovingAverage(decay=0.9999) 82 | ema_vars = tf.trainable_variables() + tf.get_collection('moving_vars') 83 | for v in tf.global_variables(): 84 | if 'moving_mean' in v.name or 'moving_variance' in v.name: 85 | ema_vars.append(v) 86 | ema_vars = list(set(ema_vars)) 87 | var_dict = ema.variables_to_restore(ema_vars) 88 | saver = tf.train.Saver(var_dict, max_to_keep=1) 89 | saver.restore(sess, checkpoint) 90 | 91 | def build_model(self, features, is_training): 92 | """Build model with input features.""" 93 | features -= tf.constant(MEAN_RGB, shape=[1, 1, 3], dtype=features.dtype) 94 | features /= tf.constant(STDDEV_RGB, shape=[1, 1, 3], dtype=features.dtype) 95 | logits, _ = efficientnet_builder.build_model( 96 | features, self.model_name, is_training) 97 | probs = tf.nn.softmax(logits) 98 | probs = tf.squeeze(probs) 99 | return probs 100 | 101 | def build_dataset(self, filenames, labels, is_training): 102 | """Build input dataset.""" 103 | filenames = tf.constant(filenames) 104 | labels = tf.constant(labels) 105 | dataset = tf.data.Dataset.from_tensor_slices((filenames, labels)) 106 | 107 | def _parse_function(filename, label): 108 | image_string = tf.read_file(filename) 109 | image_decoded = preprocessing.preprocess_image( 110 | image_string, is_training, self.image_size) 111 | image = tf.cast(image_decoded, tf.float32) 112 | return image, label 113 | 114 | dataset = dataset.map(_parse_function) 115 | dataset = dataset.batch(self.batch_size) 116 | 117 | iterator = dataset.make_one_shot_iterator() 118 | images, labels = iterator.get_next() 119 | return images, labels 120 | 121 | def run_inference(self, ckpt_dir, image_files, labels): 122 | """Build and run inference on the target images and labels.""" 123 | with tf.Graph().as_default(), tf.Session() as sess: 124 | images, labels = self.build_dataset(image_files, labels, False) 125 | probs = self.build_model(images, is_training=False) 126 | 127 | sess.run(tf.global_variables_initializer()) 128 | self.restore_model(sess, ckpt_dir) 129 | 130 | prediction_idx = [] 131 | prediction_prob = [] 132 | for _ in range(len(image_files) // self.batch_size): 133 | out_probs = sess.run(probs) 134 | idx = np.argsort(out_probs)[::-1] 135 | prediction_idx.append(idx[:5]) 136 | prediction_prob.append([out_probs[pid] for pid in idx[:5]]) 137 | 138 | # Return the top 5 predictions (idx and prob) for each image. 139 | return prediction_idx, prediction_prob 140 | 141 | 142 | def eval_example_images(model_name, ckpt_dir, image_files, labels_map_file): 143 | """Eval a list of example images. 144 | 145 | Args: 146 | model_name: str. The name of model to eval. 147 | ckpt_dir: str. Checkpoint directory path. 148 | image_files: List[str]. A list of image file paths. 149 | labels_map_file: str. The labels map file path. 150 | 151 | Returns: 152 | A tuple (pred_idx, and pred_prob), where pred_idx is the top 5 prediction 153 | index and pred_prob is the top 5 prediction probability. 154 | """ 155 | eval_ckpt_driver = EvalCkptDriver(model_name) 156 | classes = json.loads(tf.gfile.Open(labels_map_file).read()) 157 | pred_idx, pred_prob = eval_ckpt_driver.run_inference( 158 | ckpt_dir, image_files, [0] * len(image_files)) 159 | for i in range(len(image_files)): 160 | print('predicted class for image {}: '.format(image_files[i])) 161 | for j, idx in enumerate(pred_idx[i]): 162 | print(' -> top_{} ({:4.2f}%): {} '.format( 163 | j, pred_prob[i][j] * 100, classes[str(idx)])) 164 | return pred_idx, pred_prob 165 | 166 | 167 | def eval_imagenet(model_name, 168 | ckpt_dir, 169 | imagenet_eval_glob, 170 | imagenet_eval_label, 171 | num_images): 172 | """Eval ImageNet images and report top1/top5 accuracy. 173 | 174 | Args: 175 | model_name: str. The name of model to eval. 176 | ckpt_dir: str. Checkpoint directory path. 177 | imagenet_eval_glob: str. File path glob for all eval images. 178 | imagenet_eval_label: str. File path for eval label. 179 | num_images: int. Number of images to eval: -1 means eval the whole dataset. 180 | 181 | Returns: 182 | A tuple (top1, top5) for top1 and top5 accuracy. 183 | """ 184 | eval_ckpt_driver = EvalCkptDriver(model_name) 185 | imagenet_val_labels = [int(i) for i in tf.gfile.GFile(imagenet_eval_label)] 186 | imagenet_filenames = sorted(tf.gfile.Glob(imagenet_eval_glob)) 187 | if num_images < 0: 188 | num_images = len(imagenet_filenames) 189 | image_files = imagenet_filenames[:num_images] 190 | labels = imagenet_val_labels[:num_images] 191 | 192 | pred_idx, _ = eval_ckpt_driver.run_inference(ckpt_dir, image_files, labels) 193 | top1_cnt, top5_cnt = 0.0, 0.0 194 | for i, label in enumerate(labels): 195 | top1_cnt += label in pred_idx[i][:1] 196 | top5_cnt += label in pred_idx[i][:5] 197 | if i % 100 == 0: 198 | print('Step {}: top1_acc = {:4.2f}% top5_acc = {:4.2f}%'.format( 199 | i, 100 * top1_cnt / (i + 1), 100 * top5_cnt / (i + 1))) 200 | sys.stdout.flush() 201 | top1, top5 = 100 * top1_cnt / num_images, 100 * top5_cnt / num_images 202 | print('Final: top1_acc = {:4.2f}% top5_acc = {:4.2f}%'.format(top1, top5)) 203 | return top1, top5 204 | 205 | 206 | def main(unused_argv): 207 | tf.logging.set_verbosity(tf.logging.ERROR) 208 | if FLAGS.runmode == 'examples': 209 | # Run inference for an example image. 210 | eval_example_images(FLAGS.model_name, FLAGS.ckpt_dir, [FLAGS.example_img], 211 | FLAGS.labels_map_file) 212 | elif FLAGS.runmode == 'imagenet': 213 | # Run inference for imagenet. 214 | eval_imagenet(FLAGS.model_name, FLAGS.ckpt_dir, FLAGS.imagenet_eval_glob, 215 | FLAGS.imagenet_eval_label, FLAGS.num_images) 216 | else: 217 | print('must specify runmode: examples or imagenet') 218 | 219 | 220 | if __name__ == '__main__': 221 | app.run(main) 222 | -------------------------------------------------------------------------------- /convert_tf_to_pt/original_tf/preprocessing.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """ImageNet preprocessing.""" 16 | from __future__ import absolute_import 17 | from __future__ import division 18 | from __future__ import print_function 19 | 20 | import tensorflow as tf 21 | 22 | IMAGE_SIZE = 224 23 | CROP_PADDING = 32 24 | 25 | 26 | def distorted_bounding_box_crop(image_bytes, 27 | bbox, 28 | min_object_covered=0.1, 29 | aspect_ratio_range=(0.75, 1.33), 30 | area_range=(0.05, 1.0), 31 | max_attempts=100, 32 | scope=None): 33 | """Generates cropped_image using one of the bboxes randomly distorted. 34 | 35 | See `tf.image.sample_distorted_bounding_box` for more documentation. 36 | 37 | Args: 38 | image_bytes: `Tensor` of binary image data. 39 | bbox: `Tensor` of bounding boxes arranged `[1, num_boxes, coords]` 40 | where each coordinate is [0, 1) and the coordinates are arranged 41 | as `[ymin, xmin, ymax, xmax]`. If num_boxes is 0 then use the whole 42 | image. 43 | min_object_covered: An optional `float`. Defaults to `0.1`. The cropped 44 | area of the image must contain at least this fraction of any bounding 45 | box supplied. 46 | aspect_ratio_range: An optional list of `float`s. The cropped area of the 47 | image must have an aspect ratio = width / height within this range. 48 | area_range: An optional list of `float`s. The cropped area of the image 49 | must contain a fraction of the supplied image within in this range. 50 | max_attempts: An optional `int`. Number of attempts at generating a cropped 51 | region of the image of the specified constraints. After `max_attempts` 52 | failures, return the entire image. 53 | scope: Optional `str` for name scope. 54 | Returns: 55 | cropped image `Tensor` 56 | """ 57 | with tf.name_scope(scope, 'distorted_bounding_box_crop', [image_bytes, bbox]): 58 | shape = tf.image.extract_jpeg_shape(image_bytes) 59 | sample_distorted_bounding_box = tf.image.sample_distorted_bounding_box( 60 | shape, 61 | bounding_boxes=bbox, 62 | min_object_covered=min_object_covered, 63 | aspect_ratio_range=aspect_ratio_range, 64 | area_range=area_range, 65 | max_attempts=max_attempts, 66 | use_image_if_no_bounding_boxes=True) 67 | bbox_begin, bbox_size, _ = sample_distorted_bounding_box 68 | 69 | # Crop the image to the specified bounding box. 70 | offset_y, offset_x, _ = tf.unstack(bbox_begin) 71 | target_height, target_width, _ = tf.unstack(bbox_size) 72 | crop_window = tf.stack([offset_y, offset_x, target_height, target_width]) 73 | image = tf.image.decode_and_crop_jpeg(image_bytes, crop_window, channels=3) 74 | 75 | return image 76 | 77 | 78 | def _at_least_x_are_equal(a, b, x): 79 | """At least `x` of `a` and `b` `Tensors` are equal.""" 80 | match = tf.equal(a, b) 81 | match = tf.cast(match, tf.int32) 82 | return tf.greater_equal(tf.reduce_sum(match), x) 83 | 84 | 85 | def _decode_and_random_crop(image_bytes, image_size): 86 | """Make a random crop of image_size.""" 87 | bbox = tf.constant([0.0, 0.0, 1.0, 1.0], dtype=tf.float32, shape=[1, 1, 4]) 88 | image = distorted_bounding_box_crop( 89 | image_bytes, 90 | bbox, 91 | min_object_covered=0.1, 92 | aspect_ratio_range=(3. / 4, 4. / 3.), 93 | area_range=(0.08, 1.0), 94 | max_attempts=10, 95 | scope=None) 96 | original_shape = tf.image.extract_jpeg_shape(image_bytes) 97 | bad = _at_least_x_are_equal(original_shape, tf.shape(image), 3) 98 | 99 | image = tf.cond( 100 | bad, 101 | lambda: _decode_and_center_crop(image_bytes, image_size), 102 | lambda: tf.image.resize_bicubic([image], # pylint: disable=g-long-lambda 103 | [image_size, image_size])[0]) 104 | 105 | return image 106 | 107 | 108 | def _decode_and_center_crop(image_bytes, image_size): 109 | """Crops to center of image with padding then scales image_size.""" 110 | shape = tf.image.extract_jpeg_shape(image_bytes) 111 | image_height = shape[0] 112 | image_width = shape[1] 113 | 114 | padded_center_crop_size = tf.cast( 115 | ((image_size / (image_size + CROP_PADDING)) * 116 | tf.cast(tf.minimum(image_height, image_width), tf.float32)), 117 | tf.int32) 118 | 119 | offset_height = ((image_height - padded_center_crop_size) + 1) // 2 120 | offset_width = ((image_width - padded_center_crop_size) + 1) // 2 121 | crop_window = tf.stack([offset_height, offset_width, 122 | padded_center_crop_size, padded_center_crop_size]) 123 | image = tf.image.decode_and_crop_jpeg(image_bytes, crop_window, channels=3) 124 | image = tf.image.resize_bicubic([image], [image_size, image_size])[0] 125 | 126 | return image 127 | 128 | 129 | def _flip(image): 130 | """Random horizontal image flip.""" 131 | image = tf.image.random_flip_left_right(image) 132 | return image 133 | 134 | 135 | def preprocess_for_train(image_bytes, use_bfloat16, image_size=IMAGE_SIZE): 136 | """Preprocesses the given image for evaluation. 137 | 138 | Args: 139 | image_bytes: `Tensor` representing an image binary of arbitrary size. 140 | use_bfloat16: `bool` for whether to use bfloat16. 141 | image_size: image size. 142 | 143 | Returns: 144 | A preprocessed image `Tensor`. 145 | """ 146 | image = _decode_and_random_crop(image_bytes, image_size) 147 | image = _flip(image) 148 | image = tf.reshape(image, [image_size, image_size, 3]) 149 | image = tf.image.convert_image_dtype( 150 | image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32) 151 | return image 152 | 153 | 154 | def preprocess_for_eval(image_bytes, use_bfloat16, image_size=IMAGE_SIZE): 155 | """Preprocesses the given image for evaluation. 156 | 157 | Args: 158 | image_bytes: `Tensor` representing an image binary of arbitrary size. 159 | use_bfloat16: `bool` for whether to use bfloat16. 160 | image_size: image size. 161 | 162 | Returns: 163 | A preprocessed image `Tensor`. 164 | """ 165 | image = _decode_and_center_crop(image_bytes, image_size) 166 | image = tf.reshape(image, [image_size, image_size, 3]) 167 | image = tf.image.convert_image_dtype( 168 | image, dtype=tf.bfloat16 if use_bfloat16 else tf.float32) 169 | return image 170 | 171 | 172 | def preprocess_image(image_bytes, 173 | is_training=False, 174 | use_bfloat16=False, 175 | image_size=IMAGE_SIZE): 176 | """Preprocesses the given image. 177 | 178 | Args: 179 | image_bytes: `Tensor` representing an image binary of arbitrary size. 180 | is_training: `bool` for whether the preprocessing is for training. 181 | use_bfloat16: `bool` for whether to use bfloat16. 182 | image_size: image size. 183 | 184 | Returns: 185 | A preprocessed image `Tensor` with value range of [0, 255]. 186 | """ 187 | if is_training: 188 | return preprocess_for_train(image_bytes, use_bfloat16, image_size) 189 | else: 190 | return preprocess_for_eval(image_bytes, use_bfloat16, image_size) 191 | -------------------------------------------------------------------------------- /convert_tf_to_pt/original_tf/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Model utilities.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import os 22 | import numpy as np 23 | import tensorflow as tf 24 | 25 | from tensorflow.contrib.tpu.python.ops import tpu_ops 26 | from tensorflow.contrib.tpu.python.tpu import tpu_function 27 | 28 | 29 | def build_learning_rate(initial_lr, 30 | global_step, 31 | steps_per_epoch=None, 32 | lr_decay_type='exponential', 33 | decay_factor=0.97, 34 | decay_epochs=2.4, 35 | total_steps=None, 36 | warmup_epochs=5): 37 | """Build learning rate.""" 38 | if lr_decay_type == 'exponential': 39 | assert steps_per_epoch is not None 40 | decay_steps = steps_per_epoch * decay_epochs 41 | lr = tf.train.exponential_decay( 42 | initial_lr, global_step, decay_steps, decay_factor, staircase=True) 43 | elif lr_decay_type == 'cosine': 44 | assert total_steps is not None 45 | lr = 0.5 * initial_lr * ( 46 | 1 + tf.cos(np.pi * tf.cast(global_step, tf.float32) / total_steps)) 47 | elif lr_decay_type == 'constant': 48 | lr = initial_lr 49 | else: 50 | assert False, 'Unknown lr_decay_type : %s' % lr_decay_type 51 | 52 | if warmup_epochs: 53 | tf.logging.info('Learning rate warmup_epochs: %d' % warmup_epochs) 54 | warmup_steps = int(warmup_epochs * steps_per_epoch) 55 | warmup_lr = ( 56 | initial_lr * tf.cast(global_step, tf.float32) / tf.cast( 57 | warmup_steps, tf.float32)) 58 | lr = tf.cond(global_step < warmup_steps, lambda: warmup_lr, lambda: lr) 59 | 60 | return lr 61 | 62 | 63 | def build_optimizer(learning_rate, 64 | optimizer_name='rmsprop', 65 | decay=0.9, 66 | epsilon=0.001, 67 | momentum=0.9): 68 | """Build optimizer.""" 69 | if optimizer_name == 'sgd': 70 | tf.logging.info('Using SGD optimizer') 71 | optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate) 72 | elif optimizer_name == 'momentum': 73 | tf.logging.info('Using Momentum optimizer') 74 | optimizer = tf.train.MomentumOptimizer( 75 | learning_rate=learning_rate, momentum=momentum) 76 | elif optimizer_name == 'rmsprop': 77 | tf.logging.info('Using RMSProp optimizer') 78 | optimizer = tf.train.RMSPropOptimizer(learning_rate, decay, momentum, 79 | epsilon) 80 | else: 81 | tf.logging.fatal('Unknown optimizer:', optimizer_name) 82 | 83 | return optimizer 84 | 85 | 86 | class TpuBatchNormalization(tf.layers.BatchNormalization): 87 | # class TpuBatchNormalization(tf.layers.BatchNormalization): 88 | """Cross replica batch normalization.""" 89 | 90 | def __init__(self, fused=False, **kwargs): 91 | if fused in (True, None): 92 | raise ValueError('TpuBatchNormalization does not support fused=True.') 93 | super(TpuBatchNormalization, self).__init__(fused=fused, **kwargs) 94 | 95 | def _cross_replica_average(self, t, num_shards_per_group): 96 | """Calculates the average value of input tensor across TPU replicas.""" 97 | num_shards = tpu_function.get_tpu_context().number_of_shards 98 | group_assignment = None 99 | if num_shards_per_group > 1: 100 | if num_shards % num_shards_per_group != 0: 101 | raise ValueError('num_shards: %d mod shards_per_group: %d, should be 0' 102 | % (num_shards, num_shards_per_group)) 103 | num_groups = num_shards // num_shards_per_group 104 | group_assignment = [[ 105 | x for x in range(num_shards) if x // num_shards_per_group == y 106 | ] for y in range(num_groups)] 107 | return tpu_ops.cross_replica_sum(t, group_assignment) / tf.cast( 108 | num_shards_per_group, t.dtype) 109 | 110 | def _moments(self, inputs, reduction_axes, keep_dims): 111 | """Compute the mean and variance: it overrides the original _moments.""" 112 | shard_mean, shard_variance = super(TpuBatchNormalization, self)._moments( 113 | inputs, reduction_axes, keep_dims=keep_dims) 114 | 115 | num_shards = tpu_function.get_tpu_context().number_of_shards or 1 116 | if num_shards <= 8: # Skip cross_replica for 2x2 or smaller slices. 117 | num_shards_per_group = 1 118 | else: 119 | num_shards_per_group = max(8, num_shards // 4) 120 | tf.logging.info('TpuBatchNormalization with num_shards_per_group %s', 121 | num_shards_per_group) 122 | if num_shards_per_group > 1: 123 | # Each group has multiple replicas: here we compute group mean/variance by 124 | # aggregating per-replica mean/variance. 125 | group_mean = self._cross_replica_average(shard_mean, num_shards_per_group) 126 | group_variance = self._cross_replica_average(shard_variance, 127 | num_shards_per_group) 128 | 129 | # Group variance needs to also include the difference between shard_mean 130 | # and group_mean. 131 | mean_distance = tf.square(group_mean - shard_mean) 132 | group_variance += self._cross_replica_average(mean_distance, 133 | num_shards_per_group) 134 | return (group_mean, group_variance) 135 | else: 136 | return (shard_mean, shard_variance) 137 | 138 | 139 | def drop_connect(inputs, is_training, drop_connect_rate): 140 | """Apply drop connect.""" 141 | if not is_training: 142 | return inputs 143 | 144 | # Compute keep_prob 145 | # TODO(tanmingxing): add support for training progress. 146 | keep_prob = 1.0 - drop_connect_rate 147 | 148 | # Compute drop_connect tensor 149 | batch_size = tf.shape(inputs)[0] 150 | random_tensor = keep_prob 151 | random_tensor += tf.random_uniform([batch_size, 1, 1, 1], dtype=inputs.dtype) 152 | binary_tensor = tf.floor(random_tensor) 153 | output = tf.div(inputs, keep_prob) * binary_tensor 154 | return output 155 | 156 | 157 | def archive_ckpt(ckpt_eval, ckpt_objective, ckpt_path): 158 | """Archive a checkpoint if the metric is better.""" 159 | ckpt_dir, ckpt_name = os.path.split(ckpt_path) 160 | 161 | saved_objective_path = os.path.join(ckpt_dir, 'best_objective.txt') 162 | saved_objective = float('-inf') 163 | if tf.gfile.Exists(saved_objective_path): 164 | with tf.gfile.GFile(saved_objective_path, 'r') as f: 165 | saved_objective = float(f.read()) 166 | if saved_objective > ckpt_objective: 167 | tf.logging.info('Ckpt %s is worse than %s', ckpt_objective, saved_objective) 168 | return False 169 | 170 | filenames = tf.gfile.Glob(ckpt_path + '.*') 171 | if filenames is None: 172 | tf.logging.info('No files to copy for checkpoint %s', ckpt_path) 173 | return False 174 | 175 | # Clear the old folder. 176 | dst_dir = os.path.join(ckpt_dir, 'archive') 177 | if tf.gfile.Exists(dst_dir): 178 | tf.gfile.DeleteRecursively(dst_dir) 179 | tf.gfile.MakeDirs(dst_dir) 180 | 181 | # Write checkpoints. 182 | for f in filenames: 183 | dest = os.path.join(dst_dir, os.path.basename(f)) 184 | tf.gfile.Copy(f, dest, overwrite=True) 185 | ckpt_state = tf.train.generate_checkpoint_state_proto( 186 | dst_dir, 187 | model_checkpoint_path=ckpt_name, 188 | all_model_checkpoint_paths=[ckpt_name]) 189 | with tf.gfile.GFile(os.path.join(dst_dir, 'checkpoint'), 'w') as f: 190 | f.write(str(ckpt_state)) 191 | with tf.gfile.GFile(os.path.join(dst_dir, 'best_eval.txt'), 'w') as f: 192 | f.write('%s' % ckpt_eval) 193 | 194 | # Update the best objective. 195 | with tf.gfile.GFile(saved_objective_path, 'w') as f: 196 | f.write('%f' % ckpt_objective) 197 | 198 | tf.logging.info('Copying checkpoint %s to %s', ckpt_path, dst_dir) 199 | return True 200 | 201 | 202 | # TODO(hongkuny): Consolidate this as a common library cross models. 203 | class DepthwiseConv2D(tf.keras.layers.DepthwiseConv2D, tf.layers.Layer): 204 | """Wrap keras DepthwiseConv2D to tf.layers.""" 205 | 206 | pass 207 | -------------------------------------------------------------------------------- /convert_tf_to_pt/prototxt_basic.py: -------------------------------------------------------------------------------- 1 | # prototxt_basic 2 | import sys 3 | import logging 4 | logging.basicConfig(level = logging.INFO) 5 | 6 | def data(txt_file, info): 7 | txt_file.write('name: "efficientnet_pytorch2caffe"\n') 8 | txt_file.write('layer {\n') 9 | txt_file.write(' name: "data"\n') 10 | txt_file.write(' type: "Input"\n') 11 | txt_file.write(' top: "data"\n') 12 | txt_file.write(' input_param {\n') 13 | #txt_file.write(' shape: { dim: 10 dim: 3 dim: 224 dim: 224 }\n') # TODO 14 | txt_file.write(' shape: { dim: 1 dim: 3 dim: 224 dim: 224 }\n') # TODO 15 | txt_file.write(' }\n') 16 | txt_file.write('}\n') 17 | txt_file.write('\n') 18 | 19 | def SliceChannel(txt_file, info): 20 | txt_file.write('layer {\n') 21 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 22 | txt_file.write(' top: "%s"\n' % info['top']) 23 | txt_file.write(' name: "%s"\n' % info['top']) 24 | txt_file.write(' type: "SliceChannel"\n') 25 | txt_file.write('}\n') 26 | txt_file.write('\n') 27 | 28 | def L2Normalization(txt_file, info): 29 | txt_file.write('layer {\n') 30 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 31 | txt_file.write(' top: "%s"\n' % info['top']) 32 | txt_file.write(' name: "%s"\n' % info['top']) 33 | txt_file.write(' type: "L2Normalization"\n') 34 | txt_file.write('}\n') 35 | txt_file.write('\n') 36 | 37 | def DropOut(txt_file, info): 38 | txt_file.write('layer {\n') 39 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 40 | txt_file.write(' top: "%s"\n' % info['top']) 41 | txt_file.write(' name: "%s"\n' % info['top']) 42 | txt_file.write(' type: "Dropout"\n') 43 | 44 | txt_file.write(' dropout_param {\n') 45 | txt_file.write(' dropout_ratio: %s\n'% str(1.0 - float(info['attrs']['p']))) 46 | txt_file.write(' }\n') 47 | 48 | txt_file.write('}\n') 49 | txt_file.write('\n') 50 | 51 | def Convolution(txt_file, info): 52 | if info['attrs']['no_bias'] == 'True': 53 | bias_term = 'false' 54 | else: 55 | bias_term = 'true' 56 | txt_file.write('layer {\n') 57 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 58 | txt_file.write(' top: "%s"\n' % info['top']) 59 | txt_file.write(' name: "%s"\n' % info['top']) 60 | txt_file.write(' type: "Convolution"\n') 61 | txt_file.write(' convolution_param {\n') 62 | txt_file.write(' num_output: %s\n' % info['attrs']['num_filter']) 63 | txt_file.write(' kernel_size: %s\n' % info['attrs']['kernel'].split('(')[1].split(',')[0]) # TODO 64 | if 'pad' not in info['attrs']: 65 | logging.info('miss Conv_pad, make pad default: 0 ') 66 | txt_file.write(' pad: %s\n' % 0) # TODO 67 | else: 68 | #print(info['attrs']['pad']) 69 | txt_file.write(' pad: %s\n' % info['attrs']['pad'].split('(')[1].split(',')[0]) # TODO 70 | 71 | if "num_group" in info['attrs']: 72 | txt_file.write(' group: %s\n' % info['attrs']['num_group']) 73 | 74 | txt_file.write(' stride: %s\n' % info['attrs']['stride'].split('(')[1].split(',')[0]) 75 | txt_file.write(' bias_term: %s\n' % bias_term) 76 | txt_file.write(' }\n') 77 | 78 | if 'share' in info.keys() and info['share']: 79 | print("Log -> ",info['top'],info['share']) 80 | txt_file.write(' param {\n') 81 | txt_file.write(' name: "%s"\n' % info['params'][0]) 82 | txt_file.write(' }\n') 83 | txt_file.write('}\n') 84 | txt_file.write('\n') 85 | 86 | def ChannelwiseConvolution(txt_file, info): 87 | Convolution(txt_file, info) 88 | 89 | def BatchNorm(txt_file, info): 90 | txt_file.write('layer {\n') 91 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 92 | txt_file.write(' top: "%s"\n' % info['top']) 93 | txt_file.write(' name: "%s"\n' % info['top']) 94 | txt_file.write(' type: "BatchNorm"\n') 95 | txt_file.write(' batch_norm_param {\n') 96 | txt_file.write(' use_global_stats: true\n') # TODO 97 | #txt_file.write(' moving_average_fraction: 0.9\n') # TODO 98 | txt_file.write(' eps: %s\n' %info['attrs']["eps"])# TODO 99 | 100 | txt_file.write(' }\n') 101 | txt_file.write('}\n') 102 | # if info['fix_gamma'] is "False": # TODO 103 | txt_file.write('layer {\n') 104 | txt_file.write(' bottom: "%s"\n' % info['top']) 105 | txt_file.write(' top: "%s"\n' % info['top']) 106 | txt_file.write(' name: "%s_scale"\n' % info['top']) 107 | txt_file.write(' type: "Scale"\n') 108 | txt_file.write(' scale_param { bias_term: true }\n') 109 | txt_file.write('}\n') 110 | txt_file.write('\n') 111 | pass 112 | 113 | 114 | def Activation(txt_file, info): 115 | txt_file.write('layer {\n') 116 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 117 | txt_file.write(' top: "%s"\n' % info['top']) 118 | txt_file.write(' name: "%s"\n' % info['top']) 119 | if info['attrs']['act_type'] == 'sigmoid': # TODO 120 | txt_file.write(' type: "Sigmoid"\n') 121 | elif info['attrs']['act_type'] == 'relu': 122 | txt_file.write(' type: "ReLU" \n') 123 | elif info['attrs']['act_type'] == 'swish': 124 | txt_file.write(' type: "Swish" \n') 125 | else: 126 | logging.info("Unknown avtivate_function %s" % info['attrs']['act_type']) 127 | 128 | txt_file.write('}\n') 129 | txt_file.write('\n') 130 | pass 131 | 132 | def Concat(txt_file, info): 133 | #for dense net for tensorrt 134 | ''' 135 | if "concat" in info['bottom'][0]: 136 | txt_file.write('layer {\n') 137 | txt_file.write(' bottom: "%s"\n' % (info['bottom'][0])) 138 | txt_file.write(' top: "%s"\n' % (info['bottom'][0] + "_Copy")) 139 | txt_file.write(' name: "%s"\n' % (info['bottom'][0] + "_Copy")) 140 | txt_file.write(' type: "Pooling"\n') 141 | txt_file.write(' pooling_param {\n') 142 | txt_file.write(' pool: MAX\n') 143 | txt_file.write(' kernel_size: 1\n') 144 | txt_file.write(' stride: 1\n') 145 | txt_file.write(' pad: 0\n') 146 | txt_file.write(' }\n') 147 | txt_file.write('}\n') 148 | txt_file.write('\n') 149 | ''' 150 | txt_file.write('layer {\n') 151 | txt_file.write(' name: "%s"\n' % info['top']) 152 | txt_file.write(' type: "Concat"\n') 153 | if len(info['bottom']) > 2: 154 | for bottom_i in info['bottom']: 155 | txt_file.write(' bottom: "%s"\n' % bottom_i) 156 | else: 157 | if "concat" in info['bottom'][0]: 158 | txt_file.write(' bottom: "%s"\n' % (info['bottom'][0] + "")) 159 | else: 160 | txt_file.write(' bottom: "%s"\n' % (info['bottom'][0])) 161 | txt_file.write(' bottom: "%s"\n' % info['bottom'][1]) 162 | 163 | txt_file.write(' top: "%s"\n' % (info['top'])) 164 | txt_file.write('}\n') 165 | txt_file.write('\n') 166 | 167 | ''' 168 | #using slice instead for int8 quantitation 169 | txt_file.write('layer {\n') 170 | txt_file.write(' name: "%s"\n' % (info['top'] + "_Copy")) 171 | txt_file.write(' type: "Slice"\n') 172 | txt_file.write(' bottom: "%s"\n' % (info['top'])) 173 | txt_file.write(' top: "%s"\n' % (info['top'] + "_Copy")) 174 | txt_file.write(' slice_param {\n') 175 | txt_file.write(' axis: 0\n') 176 | txt_file.write(' }\n') 177 | txt_file.write('}\n') 178 | txt_file.write('\n') 179 | ''' 180 | #using pooling instead for int8 quantitation 181 | ''' 182 | txt_file.write('layer {\n') 183 | txt_file.write(' bottom: "%s"\n' % (info['top'])) 184 | txt_file.write(' top: "%s"\n' % (info['top'] + "_Copy")) 185 | txt_file.write(' name: "%s"\n' % (info['top'] + "_Copy")) 186 | txt_file.write(' type: "Pooling"\n') 187 | txt_file.write(' pooling_param {\n') 188 | txt_file.write(' pool: MAX\n') 189 | txt_file.write(' kernel_size: 1\n') 190 | txt_file.write(' stride: 1\n') 191 | txt_file.write(' pad: 0\n') 192 | txt_file.write(' }\n') 193 | txt_file.write('}\n') 194 | txt_file.write('\n') 195 | ''' 196 | 197 | pass 198 | 199 | def ElementWiseSum(txt_file, info): 200 | txt_file.write('layer {\n') 201 | txt_file.write(' name: "%s"\n' % info['top']) 202 | txt_file.write(' type: "Eltwise"\n') 203 | for bottom_i in info['bottom']: 204 | txt_file.write(' bottom: "%s"\n' % bottom_i) 205 | txt_file.write(' top: "%s"\n' % info['top']) 206 | txt_file.write(' eltwise_param {\n') 207 | txt_file.write(' operation: SUM\n') 208 | txt_file.write(' }\n') 209 | txt_file.write('}\n') 210 | txt_file.write('\n') 211 | pass 212 | 213 | def Pooling(txt_file, info): 214 | pool_type = 'AVE' if info['attrs']['pool_type'] == 'avg' else 'MAX' 215 | txt_file.write('layer {\n') 216 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 217 | txt_file.write(' top: "%s"\n' % info['top']) 218 | txt_file.write(' name: "%s"\n' % info['top']) 219 | txt_file.write(' type: "Pooling"\n') 220 | txt_file.write(' pooling_param {\n') 221 | txt_file.write(' pool: %s\n' % pool_type) # TODO 222 | txt_file.write(' kernel_size: %s\n' % info['attrs']['kernel'].split('(')[1].split(',')[0]) 223 | if 'global_pool' not in info['attrs'] or info['attrs']['global_pool'] == 'False': 224 | txt_file.write(' stride: %s\n' % info['attrs']['stride'].split('(')[1].split(',')[0]) 225 | txt_file.write(' pad: %s\n' % info['attrs']['pad'].split('(')[1].split(',')[0]) 226 | else: 227 | txt_file.write(' global_pooling: True\n') 228 | 229 | if 'pooling_convention' in info['attrs']: 230 | if info['attrs']['pooling_convention'] == "valid": 231 | #must be care for 232 | txt_file.write(' round_mode: FLOOR\n') #used for caffe model 233 | #txt_file.write(' torch_pooling: true\n') #used for tensorrt caffe parser 234 | 235 | txt_file.write(' }\n') 236 | txt_file.write('}\n') 237 | txt_file.write('\n') 238 | pass 239 | 240 | def Pooling_global(txt_file, info): 241 | if 'global_pool' not in info['attrs']: 242 | Pooling(txt_file, info) 243 | return 244 | 245 | if info['attrs']['global_pool'] == 'False': 246 | Pooling(txt_file, info) 247 | 248 | elif info['attrs']['global_pool'] == 'True': 249 | pool_type = 'AVE' if info['attrs']['pool_type'] == 'avg' else 'MAX' 250 | txt_file.write('layer {\n') 251 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 252 | txt_file.write(' top: "%s"\n' % info['top']) 253 | txt_file.write(' name: "%s"\n' % info['top']) 254 | txt_file.write(' type: "Pooling"\n') 255 | txt_file.write(' pooling_param {\n') 256 | txt_file.write(' pool: %s\n' % pool_type) 257 | txt_file.write(' global_pooling: true\n') 258 | txt_file.write(' }\n') 259 | txt_file.write('}\n') 260 | txt_file.write('\n') 261 | pass 262 | 263 | def FullyConnected(txt_file, info): 264 | txt_file.write('layer {\n') 265 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 266 | txt_file.write(' top: "%s"\n' % info['top']) 267 | txt_file.write(' name: "%s"\n' % info['top']) 268 | txt_file.write(' type: "InnerProduct"\n') 269 | txt_file.write(' inner_product_param {\n') 270 | txt_file.write(' num_output: %s\n' % info['attrs']['num_hidden']) 271 | txt_file.write(' }\n') 272 | txt_file.write('}\n') 273 | txt_file.write('\n') 274 | pass 275 | 276 | def Flatten(txt_file, info): 277 | txt_file.write('layer {\n') 278 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 279 | txt_file.write(' top: "%s"\n' % info['top']) 280 | txt_file.write(' name: "%s"\n' % info['top']) 281 | txt_file.write(' type: "Flatten"\n') 282 | txt_file.write('}\n') 283 | txt_file.write('\n') 284 | 285 | def SoftmaxOutput(txt_file, info): 286 | txt_file.write('layer {\n') 287 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 288 | txt_file.write(' top: "%s"\n' % info['top']) 289 | txt_file.write(' name: "%s"\n' % info['top']) 290 | txt_file.write(' type: "Softmax"\n') 291 | txt_file.write('}\n') 292 | txt_file.write('\n') 293 | 294 | def Reshape(txt_file,info): 295 | txt_file.write('layer {\n') 296 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 297 | txt_file.write(' top: "%s"\n' % info['top']) 298 | txt_file.write(' name: "%s"\n' % info['top']) 299 | txt_file.write(' type: "Reshape"\n') 300 | txt_file.write(' reshape_param {\n') 301 | txt_file.write(' shape {\n') 302 | for dim in info['attrs']['shape'].split('(')[1].split(')')[0].split(','): 303 | txt_file.write(' dim: %s\n' % dim) 304 | #txt_file.write(' dim: %s\n' % info['attrs']['shape'].split('(')[1].split(',')[1]) 305 | #txt_file.write(' dim: %s\n' % info['attrs']['shape'].split('(')[1].split(',')[2]) 306 | #txt_file.write(' dim: %s\n' % info['attrs']['shape'].split(')')[0].split(',')[3]) 307 | txt_file.write(' }\n') 308 | txt_file.write(' axis: 1 \n') 309 | txt_file.write(' }\n') 310 | txt_file.write('}\n') 311 | txt_file.write('\n') 312 | 313 | 314 | def broadcast_mul(txt_file,info): 315 | txt_file.write('layer {\n') 316 | txt_file.write(' name: "%s"\n' % info['top']) 317 | txt_file.write(' type: "BroadcastMul"\n') 318 | for bottom_i in info['bottom']: 319 | txt_file.write(' bottom: "%s"\n' % bottom_i) 320 | txt_file.write(' top: "%s"\n' % info['top']) 321 | txt_file.write('}\n') 322 | txt_file.write('\n') 323 | 324 | def broadcast_add(txt_file,info): 325 | txt_file.write('layer {\n') 326 | txt_file.write(' name: "%s"\n' % info['top']) 327 | txt_file.write(' type: "Eltwise"\n') 328 | for bottom_i in info['bottom']: 329 | txt_file.write(' bottom: "%s"\n' % bottom_i) 330 | txt_file.write(' top: "%s"\n' % info['top']) 331 | txt_file.write(' eltwise_param {\n') 332 | txt_file.write(' operation: SUM\n') 333 | txt_file.write(' }\n') 334 | txt_file.write('}\n') 335 | txt_file.write('\n') 336 | 337 | def mulscalar(txt_file,info): 338 | ''' 339 | txt_file.write('layer {\n') 340 | txt_file.write(' top: "%s"\n' % (info['top'] + "_second")) 341 | txt_file.write(' name: "%s"\n' % (info['top'] + "_second")) 342 | txt_file.write(' type: "DummyData"\n') 343 | txt_file.write(' dummy_data_param {\n') 344 | txt_file.write(' data_filler {\n') 345 | txt_file.write(' type: "constant"\n') 346 | txt_file.write(' value: %s\n' % info['attrs']['scalar']) 347 | txt_file.write(' }\n') 348 | txt_file.write(' shape: { dim: 1 dim: 3 dim: 224 dim: 224 }\n') # TODO 349 | txt_file.write(' }\n') 350 | txt_file.write('}\n') 351 | 352 | txt_file.write('layer {\n') 353 | txt_file.write(' name: "%s"\n' % info['top']) 354 | txt_file.write(' type: "Eltwise"\n') 355 | for bottom_i in info['bottom']: 356 | txt_file.write(' bottom: "%s"\n' % bottom_i) 357 | txt_file.write(' bottom: "%s"\n' % (info['top'] + "_second")) 358 | txt_file.write(' top: "%s"\n' % info['top']) 359 | txt_file.write(' eltwise_param {\n') 360 | txt_file.write(' operation: PROD\n') 361 | txt_file.write(' }\n') 362 | txt_file.write('}\n') 363 | txt_file.write('\n') 364 | ''' 365 | txt_file.write('layer {\n') 366 | txt_file.write(' bottom: "%s"\n' % info['bottom'][0]) 367 | txt_file.write(' top: "%s"\n' % info['top']) 368 | txt_file.write(' name: "%s"\n' % info['top']) 369 | txt_file.write(' type: "Scale"\n') 370 | txt_file.write(' scale_param {\n') 371 | txt_file.write(' bias_term: false\n') 372 | txt_file.write(' filler {\n') 373 | txt_file.write(' type: "constant"\n') 374 | txt_file.write(' value: 0.017\n') 375 | txt_file.write(' }\n') 376 | txt_file.write(' }\n') 377 | txt_file.write('}\n') 378 | txt_file.write('\n') 379 | pass 380 | # ---------------------------------------------------------------- 381 | def write_node(txt_file, info): 382 | if 'label' in info['name']: 383 | return 384 | if info['op'] == 'null' and info['name'] == 'data': 385 | data(txt_file, info) 386 | elif info['op'] == 'Convolution': 387 | Convolution(txt_file, info) 388 | elif info['op'] == 'ChannelwiseConvolution': 389 | ChannelwiseConvolution(txt_file, info) 390 | elif info['op'] == 'BatchNorm': 391 | BatchNorm(txt_file, info) 392 | elif info['op'] == 'Activation': 393 | Activation(txt_file, info) 394 | # elif info['op'] == 'ElementWiseSum': 395 | elif info['op'] == 'elemwise_add': 396 | ElementWiseSum(txt_file, info) 397 | elif info['op'] == '_Plus': 398 | ElementWiseSum(txt_file, info) 399 | elif info['op'] == 'Concat': 400 | Concat(txt_file, info) 401 | elif info['op'] == 'Pooling': 402 | #Pooling(txt_file, info) 403 | Pooling_global(txt_file, info) 404 | elif info['op'] == 'Flatten': 405 | Flatten(txt_file, info) 406 | elif info['op'] == 'FullyConnected': 407 | FullyConnected(txt_file, info) 408 | elif info['op'] == 'SoftmaxActivation' or info['op'] == 'SoftmaxOutput': 409 | SoftmaxOutput(txt_file, info) 410 | ### 411 | elif info['op'] == 'Cast': 412 | Cast(txt_file, info) 413 | elif info['op'] == 'SliceChannel': 414 | SliceChannel(txt_file, info) 415 | elif info['op'] == 'L2Normalization': 416 | L2Normalization(txt_file, info) 417 | elif info['op'] == 'Reshape': 418 | #Reshape(txt_file,info) 419 | Flatten(txt_file,info) 420 | elif info['op'] == 'Dropout': 421 | DropOut(txt_file,info) 422 | elif info['op'] == "broadcast_add": 423 | broadcast_add(txt_file,info) 424 | elif info['op'] == '_mul_scalar': 425 | mulscalar(txt_file,info) 426 | else: 427 | #logging.warn("Unknown mxnet op: %s" %info['op']) 428 | sys.exit("Warning! Unknown mxnet op:{}".format(info['op'])) 429 | 430 | 431 | 432 | 433 | 434 | -------------------------------------------------------------------------------- /convert_tf_to_pt/pytorch2caffe.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import torchvision 4 | 5 | import os 6 | import sys 7 | from model_caffe import * 8 | 9 | if not len(sys.argv) == 2: 10 | sys.exit("too less argc : model name is needed! -> efficientnet-b0/efficientnet-b1/efficientnet-b2/efficientnet-b3") 11 | 12 | model_name = sys.argv[1] 13 | model = get_from_name(model_name) #from_pretrained("efficientnet-b0")#torch.load("efficientnet-b0") 14 | model.cpu() 15 | model.eval() 16 | print(model) 17 | #print(model.state_dict()) 18 | #print(m) 19 | 20 | #input_var = Variable(torch.rand(1, 3, 224, 224)) 21 | #output_var = model(input_var) 22 | #print(output_var) 23 | #output_dir = './' 24 | # plot graph to png 25 | #plot_graph(output_var, os.path.join(output_dir, 'inception_v3.dot')) 26 | #pytorch2caffe(input_var, output_var, 'efficientnet-b0.prototxt','efficientnet-b0.caffemodel') 27 | -------------------------------------------------------------------------------- /convert_tf_to_pt/pytorch2caffe_model.py: -------------------------------------------------------------------------------- 1 | 2 | import sys 3 | import torch 4 | import caffe 5 | import numpy as np 6 | from torch.utils import model_zoo 7 | 8 | if not len(sys.argv) == 2: 9 | sys.exit("too less argc : model name is needed! -> efficientnet-b0/efficientnet-b1/efficientnet-b2/efficientnet-b3") 10 | 11 | model_name = sys.argv[1] 12 | 13 | dict = torch.load("../pretrained_pytorch/" + model_name + ".pth") 14 | model = caffe.Net("caffemodel/" + model_name + ".prototxt",caffe.TEST) 15 | 16 | for kv in model.params: 17 | print kv 18 | for k,v in dict.items(): 19 | print k,v.numpy().shape 20 | 21 | for k,v in dict.items(): 22 | #batch normalization 23 | #print("PARAMETERS -> ",k) 24 | if "bn" in k: 25 | #print("Batch_Normalization -->> ",k) 26 | if ".weight" in k: 27 | caffe_key = k.replace(".weight","_scale") 28 | print (k," --> ",caffe_key,model.params[caffe_key][0].data.shape," --> ",v.numpy().shape) 29 | if model.params[caffe_key][0].data.shape != v.numpy().shape: 30 | sys.exit("data shape is not equal!") 31 | model.params[caffe_key][0].data.flat = v.numpy().flat 32 | elif ".bias" in k: 33 | caffe_key = k.replace(".bias","_scale") 34 | print (k," --> ",caffe_key,model.params[caffe_key][1].data.shape," --> ",v.numpy().shape) 35 | if model.params[caffe_key][1].data.shape != v.numpy().shape: 36 | sys.exit("data shape is not equal!") 37 | model.params[caffe_key][1].data.flat = v.numpy().flat 38 | elif ".running_mean" in k: 39 | caffe_key = k.replace(".running_mean","") 40 | print (k," --> ",caffe_key,model.params[caffe_key][0].data.shape," --> ",v.numpy().shape) 41 | if model.params[caffe_key][0].data.shape != v.numpy().shape: 42 | sys.exit("data shape is not equal!") 43 | model.params[caffe_key][0].data.flat = v.numpy().flat 44 | model.params[caffe_key][2].data[...] = 1 45 | elif ".running_var" in k: 46 | caffe_key = k.replace(".running_var","") 47 | print (k," --> ",caffe_key,model.params[caffe_key][1].data.shape," --> ",v.numpy().shape) 48 | if model.params[caffe_key][1].data.shape != v.numpy().shape: 49 | sys.exit("data shape is not equal!") 50 | model.params[caffe_key][1].data.flat = v.numpy().flat 51 | model.params[caffe_key][2].data[...] = 1 52 | else: 53 | print(k," parammeters parser error!") 54 | else: 55 | #print("ConV/FC -->> ",k) 56 | if ".weight" in k: 57 | caffe_key = k.replace('.weight','') 58 | print (k," --> ",caffe_key,model.params[caffe_key][0].data.shape," --> ",v.numpy().shape) 59 | if model.params[caffe_key][0].data.shape != v.numpy().shape: 60 | sys.exit("data shape is not equal!") 61 | model.params[caffe_key][0].data.flat = v.numpy().flat 62 | elif ".bias" in k: 63 | caffe_key = k.replace('.bias','') 64 | print (k," --> ",caffe_key,model.params[caffe_key][1].data.shape," --> ",v.numpy().shape) 65 | if model.params[caffe_key][1].data.shape != v.numpy().shape: 66 | sys.exit("data shape is not equal!") 67 | model.params[caffe_key][1].data.flat = v.numpy().flat 68 | else: 69 | print(k," parammeters parser error!") 70 | 71 | model.save("caffemodel/" + model_name + ".caffemodel") 72 | #for kv in model.params: 73 | # print kv 74 | 75 | #print(type(dict)) 76 | #for k,v in dict.items(): 77 | # print k 78 | #print(dict.items()) 79 | #print dict 80 | 81 | 82 | -------------------------------------------------------------------------------- /convert_tf_to_pt/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xxworkspace/efficientnet_tf_pytorch_caffe/3e683de7e7667165e0c1cb60d3b0ca68c8ba9098/convert_tf_to_pt/test.jpg -------------------------------------------------------------------------------- /convert_tf_to_pt/test_caffe.py: -------------------------------------------------------------------------------- 1 | 2 | import sys 3 | import torch 4 | import caffe 5 | from PIL import Image 6 | from torchvision import transforms 7 | 8 | if not len(sys.argv) == 2: 9 | sys.exit("too less argc : model name is needed! -> efficientnet-b0/efficientnet-b1/efficientnet-b2/efficientnet-b3") 10 | 11 | model_name = sys.argv[1] 12 | deploy = "caffemodel/" + model_name + ".prototxt" 13 | caffemodel = "caffemodel/" + model_name + ".caffemodel" 14 | 15 | tfms = transforms.Compose([transforms.Resize((224,224)), transforms.ToTensor(), 16 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),]) 17 | img = tfms(Image.open('test.jpg')).unsqueeze(0).numpy() 18 | 19 | model = caffe.Net(deploy,caffemodel,caffe.TEST) 20 | model.blobs['data'].data[...] = img 21 | out = model.forward() 22 | prob= model.blobs["_fc"].data 23 | print prob.reshape(-1) 24 | print img.shape 25 | -------------------------------------------------------------------------------- /convert_tf_to_pt/test_pytorch.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import json 3 | import torch 4 | import numpy as np 5 | from PIL import Image 6 | from model import * 7 | from torchvision import transforms 8 | 9 | if not len(sys.argv) == 2: 10 | sys.exit("too less argc : model name is needed! -> efficientnet-b0/efficientnet-b1/efficientnet-b2/efficientnet-b3") 11 | 12 | model_name = sys.argv[1] 13 | model = get_from_pretrained("../pretrained_pytorch/" + model_name + ".pth") 14 | 15 | tfms = transforms.Compose([transforms.Resize((224,224)), transforms.ToTensor(), 16 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),]) 17 | img = tfms(Image.open('test.jpg')).unsqueeze(0) 18 | print(img.shape) # torch.Size([1, 3, 224, 224]) 19 | print(img.dtype) 20 | 21 | model.eval() 22 | with torch.no_grad(): 23 | outputs = model(img) 24 | 25 | print(outputs.reshape(-1).astype(np.float32)) 26 | print(outputs.shape) 27 | #print(type(outputs.numpy())) 28 | 29 | #dict = torch.load("../pretrained_pytorch/efficientnet-b0.pth") 30 | #print(dict["_conv_stem.weight"]) -------------------------------------------------------------------------------- /convert_tf_to_pt/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | This file contains helper functions for building the model and for loading model parameters. 3 | These helper functions are built to mirror those in the official TensorFlow implementation. 4 | """ 5 | import re 6 | import math 7 | import collections 8 | import torch 9 | from torch import nn 10 | from torch.nn import functional as F 11 | from torch.utils import model_zoo 12 | 13 | 14 | ######################################################################## 15 | ############### HELPERS FUNCTIONS FOR MODEL ARCHITECTURE ############### 16 | ######################################################################## 17 | 18 | 19 | # Parameters for the entire model (stem, all blocks, and head) 20 | GlobalParams = collections.namedtuple('GlobalParams', [ 21 | 'batch_norm_momentum', 'batch_norm_epsilon', 'dropout_rate', 22 | 'num_classes', 'width_coefficient', 'depth_coefficient', 23 | 'depth_divisor', 'min_depth', 'drop_connect_rate',]) 24 | 25 | 26 | # Parameters for an individual model block 27 | BlockArgs = collections.namedtuple('BlockArgs', [ 28 | 'kernel_size', 'num_repeat', 'input_filters', 'output_filters', 29 | 'expand_ratio', 'id_skip', 'stride', 'se_ratio']) 30 | 31 | 32 | # Change namedtuple defaults 33 | GlobalParams.__new__.__defaults__ = (None,) * len(GlobalParams._fields) 34 | BlockArgs.__new__.__defaults__ = (None,) * len(BlockArgs._fields) 35 | 36 | class Swish(torch.nn.Module): 37 | def __init__(self): 38 | super(Swish,self).__init__() 39 | def forward(self,x): 40 | return x*F.sigmoid(x) 41 | 42 | class AddTensor(torch.nn.Module): 43 | def __init__(self): 44 | super(AddTensor,self).__init__() 45 | def forward(self,x0,x): 46 | return x0 + x 47 | 48 | class BroadcastMul(torch.nn.Module): 49 | def __init__(self): 50 | super(BroadcastMul,self).__init__() 51 | def forward(self,x,scale): 52 | return x*scale 53 | 54 | def round_filters(filters, global_params): 55 | """ Calculate and round number of filters based on depth multiplier. """ 56 | multiplier = global_params.width_coefficient 57 | if not multiplier: 58 | return filters 59 | divisor = global_params.depth_divisor 60 | min_depth = global_params.min_depth 61 | filters *= multiplier 62 | min_depth = min_depth or divisor 63 | new_filters = max(min_depth, int(filters + divisor / 2) // divisor * divisor) 64 | if new_filters < 0.9 * filters: # prevent rounding by more than 10% 65 | new_filters += divisor 66 | return int(new_filters) 67 | 68 | def round_repeats(repeats, global_params): 69 | """ Round number of filters based on depth multiplier. """ 70 | multiplier = global_params.depth_coefficient 71 | if not multiplier: 72 | return repeats 73 | return int(math.ceil(multiplier * repeats)) 74 | 75 | def getConv2d(in_size, in_channels, out_size, out_channels, kernel_size, stride=1, dilation=1, groups=1, bias=True): 76 | pad = max((out_size - 1) * stride + (kernel_size - 1) * dilation + 1 - in_size, 0) 77 | pad = int(pad - pad//2) 78 | return torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding = pad, dilation=dilation, groups=groups, bias=bias) 79 | ######################################################################## 80 | ############## HELPERS FUNCTIONS FOR LOADING MODEL PARAMS ############## 81 | ######################################################################## 82 | 83 | 84 | def efficientnet_params(model_name): 85 | """ Map EfficientNet model name to parameter coefficients. """ 86 | params_dict = { 87 | # Coefficients: width,depth,res,dropout 88 | 'efficientnet-b0': (1.0, 1.0, 224, 0.2), 89 | 'efficientnet-b1': (1.0, 1.1, 240, 0.2), 90 | 'efficientnet-b2': (1.1, 1.2, 260, 0.3), 91 | 'efficientnet-b3': (1.2, 1.4, 300, 0.3), 92 | 'efficientnet-b4': (1.4, 1.8, 380, 0.4), 93 | 'efficientnet-b5': (1.6, 2.2, 456, 0.4), 94 | 'efficientnet-b6': (1.8, 2.6, 528, 0.5), 95 | 'efficientnet-b7': (2.0, 3.1, 600, 0.5), 96 | } 97 | return params_dict[model_name] 98 | 99 | 100 | class BlockDecoder(object): 101 | """ Block Decoder for readability, straight from the official TensorFlow repository """ 102 | 103 | @staticmethod 104 | def _decode_block_string(block_string): 105 | """ Gets a block through a string notation of arguments. """ 106 | assert isinstance(block_string, str) 107 | 108 | ops = block_string.split('_') 109 | options = {} 110 | for op in ops: 111 | splits = re.split(r'(\d.*)', op) 112 | if len(splits) >= 2: 113 | key, value = splits[:2] 114 | options[key] = value 115 | 116 | # Check stride 117 | assert (('s' in options and len(options['s']) == 1) or 118 | (len(options['s']) == 2 and options['s'][0] == options['s'][1])) 119 | 120 | return BlockArgs( 121 | kernel_size=int(options['k']), 122 | num_repeat=int(options['r']), 123 | input_filters=int(options['i']), 124 | output_filters=int(options['o']), 125 | expand_ratio=int(options['e']), 126 | id_skip=('noskip' not in block_string), 127 | se_ratio=float(options['se']) if 'se' in options else None, 128 | stride=[int(options['s'][0])]) 129 | 130 | @staticmethod 131 | def _encode_block_string(block): 132 | """Encodes a block to a string.""" 133 | args = [ 134 | 'r%d' % block.num_repeat, 135 | 'k%d' % block.kernel_size, 136 | 's%d%d' % (block.strides[0], block.strides[1]), 137 | 'e%s' % block.expand_ratio, 138 | 'i%d' % block.input_filters, 139 | 'o%d' % block.output_filters 140 | ] 141 | if 0 < block.se_ratio <= 1: 142 | args.append('se%s' % block.se_ratio) 143 | if block.id_skip is False: 144 | args.append('noskip') 145 | return '_'.join(args) 146 | 147 | @staticmethod 148 | def decode(string_list): 149 | """ 150 | Decodes a list of string notations to specify blocks inside the network. 151 | 152 | :param string_list: a list of strings, each string is a notation of block 153 | :return: a list of BlockArgs namedtuples of block args 154 | """ 155 | assert isinstance(string_list, list) 156 | blocks_args = [] 157 | for block_string in string_list: 158 | blocks_args.append(BlockDecoder._decode_block_string(block_string)) 159 | return blocks_args 160 | 161 | @staticmethod 162 | def encode(blocks_args): 163 | """ 164 | Encodes a list of BlockArgs to a list of strings. 165 | 166 | :param blocks_args: a list of BlockArgs namedtuples of block args 167 | :return: a list of strings, each string is a notation of block 168 | """ 169 | block_strings = [] 170 | for block in blocks_args: 171 | block_strings.append(BlockDecoder._encode_block_string(block)) 172 | return block_strings 173 | 174 | 175 | def efficientnet(width_coefficient=None, depth_coefficient=None, 176 | dropout_rate=0.2, drop_connect_rate=0.2): 177 | """ Creates a efficientnet model. """ 178 | 179 | blocks_args = [ 180 | 'r1_k3_s11_e1_i32_o16_se0.25', 'r2_k3_s22_e6_i16_o24_se0.25', 181 | 'r2_k5_s22_e6_i24_o40_se0.25', 'r3_k3_s22_e6_i40_o80_se0.25', 182 | 'r3_k5_s11_e6_i80_o112_se0.25', 'r4_k5_s22_e6_i112_o192_se0.25', 183 | 'r1_k3_s11_e6_i192_o320_se0.25', 184 | ] 185 | blocks_args = BlockDecoder.decode(blocks_args) 186 | 187 | global_params = GlobalParams( 188 | batch_norm_momentum=0.99, 189 | batch_norm_epsilon=1e-3, 190 | dropout_rate=dropout_rate, 191 | drop_connect_rate=drop_connect_rate, 192 | # data_format='channels_last', # removed, this is always true in PyTorch 193 | num_classes=1000, 194 | width_coefficient=width_coefficient, 195 | depth_coefficient=depth_coefficient, 196 | depth_divisor=8, 197 | min_depth=None 198 | ) 199 | 200 | return blocks_args, global_params 201 | 202 | 203 | def get_model_params(model_name, override_params): 204 | """ Get the block args and global params for a given model """ 205 | if model_name.startswith('efficientnet'): 206 | w, d, _, p = efficientnet_params(model_name) 207 | # note: all models have drop connect rate = 0.2 208 | blocks_args, global_params = efficientnet(width_coefficient=w, depth_coefficient=d, dropout_rate=p) 209 | else: 210 | raise NotImplementedError('model name is not pre-defined: %s' % model_name) 211 | if override_params: 212 | # ValueError will be raised here if override_params has fields not included in global_params. 213 | global_params = global_params._replace(**override_params) 214 | return blocks_args, global_params 215 | 216 | 217 | url_map = { 218 | 'efficientnet-b0': 'http://storage.googleapis.com/public-models/efficientnet-b0-08094119.pth', 219 | 'efficientnet-b1': 'http://storage.googleapis.com/public-models/efficientnet-b1-dbc7070a.pth', 220 | 'efficientnet-b2': 'http://storage.googleapis.com/public-models/efficientnet-b2-27687264.pth', 221 | 'efficientnet-b3': 'http://storage.googleapis.com/public-models/efficientnet-b3-c8376fa2.pth', 222 | 'efficientnet-b4': 'http://storage.googleapis.com/public-models/efficientnet-b4-e116e8b3.pth', 223 | 'efficientnet-b5': 'http://storage.googleapis.com/public-models/efficientnet-b5-586e6cc6.pth', 224 | } 225 | 226 | def load_pretrained_weights(model, model_name): 227 | """ Loads pretrained weights, and downloads if loading for the first time. """ 228 | state_dict = torch.load("../pretrained_pytorch/" + model_name + ".pth")#model_zoo.load_url(url_map[model_name]) 229 | model.load_state_dict(state_dict) 230 | print('Loaded pretrained weights for {}'.format(model_name)) 231 | -------------------------------------------------------------------------------- /pretrained_tensorflow/download.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | # This script accepts a single command-line argument, which specifies which model to download. 4 | # Only the b0, b1, b2, and b3 models have been released, so your command must be one of them. 5 | 6 | # For example, to download efficientnet-b0, run: 7 | # ./download.sh efficientnet-b0 8 | # And to download efficientnet-b3, run: 9 | # ./download.sh efficientnet-b3 10 | 11 | MODEL=$1 12 | wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/${MODEL}.tar.gz 13 | tar xvf ${MODEL}.tar.gz 14 | rm ${MODEL}.tar.gz 15 | --------------------------------------------------------------------------------