├── .gitignore ├── README.md └── code ├── convert.py ├── convert_densenet.py ├── densenet_tf.py ├── resnet_tf.py ├── scale_layer_tf.py ├── test_for_fp16.py ├── test_for_fp32.py ├── test_for_keras.py ├── train.py └── train_densenet_tf.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | env/ 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | 49 | # Translations 50 | *.mo 51 | *.pot 52 | 53 | # Django stuff: 54 | *.log 55 | local_settings.py 56 | 57 | # Flask stuff: 58 | instance/ 59 | .webassets-cache 60 | 61 | # Scrapy stuff: 62 | .scrapy 63 | 64 | # Sphinx documentation 65 | docs/_build/ 66 | 67 | # PyBuilder 68 | target/ 69 | 70 | # Jupyter Notebook 71 | .ipynb_checkpoints 72 | 73 | # pyenv 74 | .python-version 75 | 76 | # celery beat schedule file 77 | celerybeat-schedule 78 | 79 | # SageMath parsed files 80 | *.sage.py 81 | 82 | # dotenv 83 | .env 84 | 85 | # virtualenv 86 | .venv 87 | venv/ 88 | ENV/ 89 | 90 | # Spyder project settings 91 | .spyderproject 92 | .spyproject 93 | 94 | # Rope project settings 95 | .ropeproject 96 | 97 | # mkdocs documentation 98 | /site 99 | 100 | # mypy 101 | .mypy_cache/ 102 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # tensorRT_for_keras 2 | 使用tensorRT来加速keras代码 3 | 4 | 5 | 6 | 这里实现`resnet` ： 7 | 8 | `keras` 实现地址为：[https://github.com/flyyufelix/cnn_finetune](https://github.com/flyyufelix/cnn_finetune) 9 | 10 | `tensorRT `使用可以参考[https://github.com/parallel-forall/code-samples/tree/master/posts/TensorRT-3.0](https://github.com/parallel-forall/code-samples/tree/master/posts/TensorRT-3.0) 11 | 12 | 13 | 14 | **相关说明：** 15 | 16 | 1. `resnet_tf.py` 为调整后的`keras` 实现，使用了`tensorflow` 中的`keras` ； 17 | 2. `train.py` 训练模型，并保存为`freeze_graph` ； 18 | 3. `convert.py` 将`freeze_graph` 转换成`tensorRT` 中的`engine` ； 19 | 4. `test_*.py` 为测试代码，`*keras` 为没有使用`tensorRT` 加速，`*fp16` 、`fp32` 为使用了`tensorRT` 加速； 20 | 21 | 22 | 23 | 为了实现`tensorRT` 对`keras` 代码加速，需进行一些修改(可参考已修改的代码，`resnet_tf.py` )： 24 | 25 | 1. `import` 库需修改： 26 | 27 | ```python 28 | from keras.models import Sequential 29 | from keras.optimizers import SGD 30 | from keras.layers import Input, Dense, Convolution2D, MaxPooling2D, AveragePooling2D, ZeroPadding2D, Dropout, Flatten, merge, Reshape, Activation 31 | from keras.layers.normalization import BatchNormalization 32 | from keras.models import Model 33 | from keras import backend as K 34 | ``` 35 | 36 | 更改为 37 | 38 | ```python 39 | from tensorflow.python.keras._impl.keras.optimizers import SGD 40 | from tensorflow.python.keras._impl.keras.preprocessing.image import ImageDataGenerator 41 | from tensorflow.python.keras._impl.keras import backend as K 42 | from tensorflow.python.keras._impl.keras.callbacks import ModelCheckpoint 43 | from tensorflow.python.keras._impl.keras.callbacks import EarlyStopping 44 | from tensorflow.python.keras._impl.keras.callbacks import LearningRateScheduler 45 | from tensorflow.python.keras._impl.keras.callbacks import TensorBoard 46 | ``` 47 | 48 | 2. `BatchNormalization` 修改： 49 | 50 | 这里不能使用 51 | 52 | ```python 53 | from tensorflow.python.keras._impl.keras.layers.normalization import BatchNormalization 54 | ``` 55 | 56 | 主要是由于`tensorRT` 只支持`tensorflow` 中的`fused batch normalization` ，因此，应使用： 57 | 58 | ```python 59 | from tensorflow.python.layers.normalization import BatchNormalization 60 | ``` 61 | 62 | **注意：这里修改了`batchnormalization` ，预先训练好的模型导不进来，需要寻找方法解决这个问题，或者重新训练模型(从0开始训练)** 63 | 64 | 65 | 66 | **使用中注意事项：** 67 | 68 | 1. `tensorboard` 使用：在`callbacks` 中添加`tensorboard` 时，其参数`histogram_freq=0` ，必须设置为0，否则会报错，参考：[https://github.com/tensorflow/tensorflow/pull/9787#issuecomment-360039499](https://github.com/tensorflow/tensorflow/pull/9787#issuecomment-360039499) 69 | 2. 转换成`freeze_graph` 时，`freeze_graph.freeze_graph` 需要指定输出层的名字`output_node_names` ，这里`output_node_names=指定名字/功能` ，如`output_node_names=fc3/Softmax` ，`fc3` 为人工指定名字，`Softmax` 为实现功能； 70 | 3. 转换成`engine` 时，`trt.utils.uff_to_trt_engine` 有个参数`max_batch_size` ，需根据显存设置，不能太大； 71 | 4. 官网帮助文档可能还没有更新，使用时需注意； 72 | 5. 屏蔽`log` 输出，可以将`log_sev` 替换成`logger_severity=trt.infer.LogSeverity.ERROR` ，详见：[http://note.youdao.com/noteshare?id=b3fdec4fc9e5861c753987c0196675ef&sub=F20AB86FBA0B4547B0FD3D7930DE4988](http://note.youdao.com/noteshare?id=b3fdec4fc9e5861c753987c0196675ef&sub=F20AB86FBA0B4547B0FD3D7930DE4988) ，也可以参考代码中使用方法； 73 | 74 | 75 | -------------------------------------------------------------------------------- /code/convert.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | # convert.py 5 | # @Author : yuanwenjin 6 | # @Mail : yfor1008@gmail.com 7 | # @Date : 2018-1-18 11:51:59 8 | # @Explanation : 转RT 9 | """ 10 | 11 | import tensorrt as trt 12 | import uff 13 | from tensorrt.parsers import uffparser 14 | 15 | if __name__ == '__main__': 16 | 17 | config = { 18 | # Training params 19 | "train_data_dir": "/home/docker/code/app/get_samples/samples_train/train", # training data 20 | "val_data_dir": "/home/docker/code/app/get_samples/samples_train/valid", # validation data 21 | "train_batch_size": 8, # training batch size 22 | "epochs": 3, # number of training epochs 23 | # "num_train_samples" : 217728, # number of training examples 24 | # "num_val_samples" : 55872, # number of test examples 25 | "num_train_samples" : 500, # number of training examples 26 | "num_val_samples" : 300, # number of test examples 27 | 28 | # Where to save models (Tensorflow + TensorRT) 29 | "graphdef_file": "./model/resnet_tf_graphdef.pb", 30 | "frozen_model_file": "./model/resnet_tf_frozen_model.pb", 31 | "snapshot_dir": "./model/snapshot", 32 | "engine_save_dir": "./model/", 33 | 34 | # Needed for TensorRT 35 | "image_dim": 480, # the image size (square images) 36 | "inference_batch_size": 256, # inference batch size 37 | "input_layer": "data", # name of the input tensor in the TF computational graph 38 | "out_layer": "fc3/Softmax", # name of the output tensorf in the TF conputational graph 39 | "output_size" : 3, # number of classes in output (3) 40 | "precision": "fp16", # desired precision (fp32, fp16) 41 | 42 | "test_image_path" : "/home/data/val/roses" 43 | } 44 | 45 | G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.INFO) 46 | 47 | # Define network parameters, including inference batch size, name & dimensionality of input/output layers 48 | INPUT_LAYERS = [config['input_layer']] 49 | OUTPUT_LAYERS = [config['out_layer']] 50 | INFERENCE_BATCH_SIZE = config['inference_batch_size'] 51 | 52 | INPUT_C = 3 53 | INPUT_H = config['image_dim'] 54 | INPUT_W = config['image_dim'] 55 | 56 | # Load your newly created Tensorflow frozen model and convert it to UFF 57 | uff_model = uff.from_tensorflow_frozen_model(config['frozen_model_file'], OUTPUT_LAYERS) 58 | 59 | # Create a UFF parser to parse the UFF file created from your TF Frozen model 60 | parser = uffparser.create_uff_parser() 61 | parser.register_input(INPUT_LAYERS[0], (INPUT_C, INPUT_H, INPUT_W), 0) 62 | parser.register_output(OUTPUT_LAYERS[0]) 63 | 64 | # Build your TensorRT inference engine 65 | if config['precision'] == 'fp32': 66 | engine = trt.utils.uff_to_trt_engine(G_LOGGER, 67 | uff_model, 68 | parser, 69 | INFERENCE_BATCH_SIZE, 70 | 1<<20, 71 | trt.infer.DataType.FLOAT 72 | ) 73 | 74 | elif config['precision'] == 'fp16': 75 | engine = trt.utils.uff_to_trt_engine(G_LOGGER, 76 | uff_model, 77 | parser, 78 | INFERENCE_BATCH_SIZE, 79 | 1<<20, 80 | trt.infer.DataType.HALF 81 | ) 82 | 83 | # Serialize TensorRT engine to a file for when you are ready to deploy your model. 84 | save_path = str(config['engine_save_dir']) + "resnet_tf_b" \ 85 | + str(INFERENCE_BATCH_SIZE) + "_"+ str(config['precision']) + ".engine" 86 | 87 | trt.utils.write_engine_to_file(save_path, engine.serialize()) 88 | 89 | print("Saved TRT engine to {}".format(save_path)) 90 | -------------------------------------------------------------------------------- /code/convert_densenet.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | # convert_densenet.py 5 | # @Author : yuanwenjin 6 | # @Mail : yfor1008@gmail.com 7 | # @Date : 2018-3-1 18:17:55 8 | # @Explanation : 转RT，densenet 9 | """ 10 | 11 | import tensorrt as trt 12 | import uff 13 | from tensorrt.parsers import uffparser 14 | 15 | if __name__ == '__main__': 16 | 17 | config = { 18 | # Where to save models (Tensorflow + TensorRT) 19 | "frozen_model_file": "./model_densenet/densenet_tf_frozen_model.pb", 20 | "engine_save_dir": "./model_densenet/", 21 | 22 | # Needed for TensorRT 23 | "image_dim": 480, # the image size (square images) 24 | "inference_batch_size": 16, # inference batch size 25 | "input_layer": "data", # name of the input tensor in the TF computational graph 26 | "out_layer": "fc6/Softmax", # name of the output tensorf in the TF conputational graph 27 | "output_size" : 3, # number of classes in output (3) 28 | "precision": "fp32", # desired precision (fp32, fp16) 29 | 30 | } 31 | 32 | G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR) 33 | # G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.INFO) 34 | # G_LOGGER = trt.infer.ConsoleLogger(0) 35 | 36 | # Define network parameters, including inference batch size, name & dimensionality of input/output layers 37 | INPUT_LAYERS = [config['input_layer']] 38 | OUTPUT_LAYERS = [config['out_layer']] 39 | INFERENCE_BATCH_SIZE = config['inference_batch_size'] 40 | 41 | INPUT_C = config['output_size'] 42 | INPUT_H = config['image_dim'] 43 | INPUT_W = config['image_dim'] 44 | 45 | # Load your newly created Tensorflow frozen model and convert it to UFF 46 | uff_model = uff.from_tensorflow_frozen_model(config['frozen_model_file'], OUTPUT_LAYERS) 47 | 48 | # Create a UFF parser to parse the UFF file created from your TF Frozen model 49 | parser = uffparser.create_uff_parser() 50 | parser.register_input(INPUT_LAYERS[0], (INPUT_C, INPUT_H, INPUT_W), 0) 51 | parser.register_output(OUTPUT_LAYERS[0]) 52 | 53 | # Build your TensorRT inference engine 54 | if config['precision'] == 'fp32': 55 | engine = trt.utils.uff_to_trt_engine(G_LOGGER, 56 | uff_model, 57 | parser, 58 | INFERENCE_BATCH_SIZE, 59 | 1<<20, 60 | trt.infer.DataType.FLOAT 61 | ) 62 | 63 | elif config['precision'] == 'fp16': 64 | engine = trt.utils.uff_to_trt_engine(G_LOGGER, 65 | uff_model, 66 | parser, 67 | INFERENCE_BATCH_SIZE, 68 | 1<<20, 69 | trt.infer.DataType.HALF 70 | ) 71 | 72 | # Serialize TensorRT engine to a file for when you are ready to deploy your model. 73 | save_path = str(config['engine_save_dir']) + "densenet_tf_b" \ 74 | + str(INFERENCE_BATCH_SIZE) + "_"+ str(config['precision']) + ".engine" 75 | 76 | trt.utils.write_engine_to_file(save_path, engine.serialize()) 77 | 78 | print("Saved TRT engine to {}".format(save_path)) -------------------------------------------------------------------------------- /code/densenet_tf.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | # densenet_tf.py 5 | # @Author : yuanwenjin 6 | # @Mail : yfor1008@gmail.com 7 | # @Date : 2018-3-1 16:23:53 8 | # @Explanation : 使用tensorflow中的keras实现，方便tensorRT加速，这里去除了dropout层 9 | """ 10 | 11 | from tensorflow.python.keras._impl.keras.layers import Input, Dense, Conv2D, MaxPooling2D, GlobalAveragePooling2D, AveragePooling2D 12 | from tensorflow.python.keras._impl.keras.layers import ZeroPadding2D, add, Activation, concatenate 13 | # from tensorflow.python.keras._impl.keras.layers.normalization import BatchNormalization 14 | from tensorflow.python.layers.normalization import BatchNormalization 15 | from tensorflow.python.keras._impl.keras.models import Model 16 | 17 | from scale_layer_tf import Scale 18 | 19 | def conv_block(tensor_x, stage, branch, nb_filter): 20 | ''' 21 | ### 说明: 22 | - Apply BatchNorm, Relu, bottleneck 1x1 Conv2D, 3x3 Conv2D 23 | 24 | ### Arguments 25 | - tensor_x: input tensor 26 | - stage: index for dense block 27 | - branch: layer index within each dense block 28 | - nb_filter: number of filters 29 | ''' 30 | eps = 1.1e-5 31 | conv_name_base = 'conv' + str(stage) + '_' + str(branch) 32 | relu_name_base = 'relu' + str(stage) + '_' + str(branch) 33 | 34 | # 1x1 Convolution (Bottleneck layer) 35 | inter_channel = nb_filter * 4 36 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name=conv_name_base + '_x1_bn')(tensor_x) 37 | tensor_x = Scale(axis=3, name=conv_name_base + '_x1_scale')(tensor_x) 38 | tensor_x = Activation('relu', name=relu_name_base + '_x1')(tensor_x) 39 | tensor_x = Conv2D(inter_channel, (1, 1), name=conv_name_base + '_x1', use_bias=False)(tensor_x) 40 | 41 | # 3x3 Convolution 42 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name=conv_name_base+'_x2_bn')(tensor_x) 43 | tensor_x = Scale(axis=3, name=conv_name_base + '_x2_scale')(tensor_x) 44 | tensor_x = Activation('relu', name=relu_name_base + '_x2')(tensor_x) 45 | tensor_x = ZeroPadding2D((1, 1), name=conv_name_base + '_x2_zeropadding')(tensor_x) 46 | tensor_x = Conv2D(nb_filter, (3, 3), name=conv_name_base + '_x2', use_bias=False)(tensor_x) 47 | 48 | return tensor_x 49 | 50 | def transition_block(tensor_x, stage, nb_filter, compression=1.0): 51 | ''' 52 | # 说明: 53 | Apply BatchNorm, 1x1 Convolution, averagePooling, optional compression 54 | # Arguments 55 | - tensor_x: input tensor 56 | - stage: index for dense block 57 | - nb_filter: number of filters 58 | - compression: calculated as 1 - reduction. Reduces the number of feature maps in the transition block. 59 | ''' 60 | 61 | eps = 1.1e-5 62 | conv_name_base = 'conv' + str(stage) + '_blk' 63 | relu_name_base = 'relu' + str(stage) + '_blk' 64 | pool_name_base = 'pool' + str(stage) 65 | 66 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name=conv_name_base+'_bn')(tensor_x) 67 | tensor_x = Scale(axis=3, name=conv_name_base+'_scale')(tensor_x) 68 | tensor_x = Activation('relu', name=relu_name_base)(tensor_x) 69 | tensor_x = Conv2D(int(nb_filter * compression), (1, 1), name=conv_name_base, use_bias=False)(tensor_x) 70 | 71 | tensor_x = AveragePooling2D((2, 2), strides=(2, 2), name=pool_name_base)(tensor_x) 72 | 73 | return tensor_x 74 | 75 | def dense_block(tensor_x, stage, nb_layers, nb_filter, growth_rate, grow_nb_filters=True): 76 | ''' 77 | # 说明: 78 | Build a dense_block where the output of each conv_block is fed to subsequent ones 79 | # Arguments 80 | - tensor_x: input tensor 81 | - stage: index for dense block 82 | - nb_layers: the number of layers of conv_block to append to the model. 83 | - nb_filter: number of filters 84 | - growth_rate: growth rate 85 | - grow_nb_filters: flag to decide to allow number of filters to grow 86 | ''' 87 | 88 | # eps = 1.1e-5 89 | concat_feat = tensor_x 90 | 91 | for i in range(nb_layers): 92 | branch = i + 1 93 | tensor_x = conv_block(concat_feat, stage, branch, growth_rate) 94 | concat_feat = concatenate([concat_feat, tensor_x], axis=3, \ 95 | name='concat_' + str(stage) + '_' + str(branch)) 96 | # concat_feat = add([concat_feat, tensor_x]) 97 | 98 | if grow_nb_filters: 99 | nb_filter += growth_rate 100 | 101 | return concat_feat, nb_filter 102 | 103 | def densenet_model(img_rows, img_cols, color_type=3, num_classes=3, 104 | nb_dense_block=4, growth_rate=None, 105 | nb_filter=None, nb_layers=None, 106 | reduction=0.5, 107 | model_size=None): 108 | ''' 109 | ### 说明: 110 | - DenseNet Model for Keras 111 | - Model Schema is based on https://github.com/flyyufelix/DenseNet-Keras 112 | - ImageNet Pretrained Weights 113 | 114 | ### Arguments: 115 | - img_rows: (int), image height 116 | - img_cols: (int), image width 117 | - color_type: (int), image channel 118 | - nb_dense_block: (int), number of dense blocks to add to end 119 | - growth_rate: (int), number of filters to add per dense block 120 | - nb_filter: (int), initial number of filters 121 | - nb_layers: (list), number of filters in each block 122 | - reduction: (float), reduction factor of transition blocks. 123 | - model_size: (string), layer number of model 124 | ### Returns 125 | - A Keras model instance. 126 | ''' 127 | eps = 1.1e-5 128 | 129 | # compute compression factor 130 | compression = 1.0 - reduction 131 | 132 | img_input = Input(shape=(img_rows, img_cols, color_type), name='data') 133 | 134 | # Initial convolution 135 | tensor_x = ZeroPadding2D((3, 3), name='conv1_zeropadding')(img_input) 136 | tensor_x = Conv2D(nb_filter, (7, 7), strides=(2, 2), name='conv1', use_bias=False)(tensor_x) 137 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name='conv1_bn')(tensor_x) 138 | tensor_x = Scale(axis=3, name='conv1_scale')(tensor_x) 139 | tensor_x = Activation('relu', name='relu1')(tensor_x) 140 | tensor_x = ZeroPadding2D((1, 1), name='pool1_zeropadding')(tensor_x) 141 | tensor_x = MaxPooling2D((3, 3), strides=(2, 2), name='pool1')(tensor_x) 142 | 143 | # Add dense blocks 144 | for block_idx in range(nb_dense_block - 1): 145 | stage = block_idx + 2 146 | tensor_x, nb_filter = dense_block(tensor_x, stage, nb_layers[block_idx], nb_filter, growth_rate) 147 | 148 | # Add transition_block 149 | tensor_x = transition_block(tensor_x, stage, nb_filter, compression=compression) 150 | nb_filter = int(nb_filter * compression) 151 | 152 | final_stage = stage + 1 153 | tensor_x, nb_filter = dense_block(tensor_x, final_stage, nb_layers[-1], nb_filter, growth_rate) 154 | 155 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name='conv' + str(final_stage) + '_blk_bn')(tensor_x) 156 | tensor_x = Scale(axis=3, name='conv' + str(final_stage) + '_blk_scale')(tensor_x) 157 | tensor_x = Activation('relu', name='relu' + str(final_stage) + '_blk')(tensor_x) 158 | 159 | x_fc = GlobalAveragePooling2D(name='pool' + str(final_stage))(tensor_x) 160 | x_fc = Dense(num_classes, activation='softmax', name='fc6')(x_fc) 161 | # x_fc = Activation('softmax', name='prob')(x_fc) 162 | 163 | model = Model(img_input, x_fc, name='densenet') 164 | 165 | return model 166 | 167 | def generate_densenet(img_rows, img_cols, color_type, 168 | reduction=0.5, 169 | model_size=None, num_classes=None): 170 | ''' 171 | ### 说明： 172 | - 生成需要的模型 173 | 174 | ### 参数： 175 | - img_rows: (int), image height 176 | - img_cols: (int), image width 177 | - color_type: (int), image channel 178 | - reduction: (float), reduction factor of transition blocks. 179 | - model_size: (int), layer number of model 180 | - num_classes: (int), classes 181 | 182 | ### 返回： 183 | - MODEL: 生成好的模型 184 | ''' 185 | 186 | # dense模型参数: 121/169/161 187 | if model_size == 121: 188 | nb_dense_block = 4 189 | growth_rate = 32 190 | nb_filter = 64 191 | nb_layers = [6, 12, 24, 16] 192 | elif model_size == 169: 193 | nb_dense_block = 4 194 | growth_rate = 32 195 | nb_filter = 64 196 | nb_layers = [6, 12, 32, 32] 197 | elif model_size == 161: 198 | nb_dense_block = 4 199 | growth_rate = 48 200 | nb_filter = 96 201 | nb_layers = [6, 12, 36, 24] 202 | 203 | model = densenet_model(img_rows=img_rows, img_cols=img_cols, color_type=color_type, num_classes=num_classes, 204 | nb_dense_block=nb_dense_block, growth_rate=growth_rate, 205 | nb_filter=nb_filter, nb_layers=nb_layers, 206 | reduction=reduction, 207 | model_size=model_size) 208 | 209 | return model 210 | 211 | if __name__ == '__main__': 212 | 213 | model = generate_densenet(480, 480, 3, model_size=121, num_classes=3) 214 | model.summary() 215 | -------------------------------------------------------------------------------- /code/resnet_tf.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | # resnet_tf.py 5 | # @Author : yuanwenjin 6 | # @Mail : yfor1008@gmail.com 7 | # @Date : 2018-1-17 10:29:01 8 | # @Explanation : 使用tensorflow中的keras实现 9 | """ 10 | 11 | from tensorflow.python.keras._impl.keras.layers import Input, Dense, Conv2D, MaxPooling2D, GlobalAveragePooling2D 12 | from tensorflow.python.keras._impl.keras.layers import ZeroPadding2D, add, Activation 13 | # from tensorflow.python.keras._impl.keras.layers.normalization import BatchNormalization 14 | from tensorflow.python.layers.normalization import BatchNormalization 15 | from tensorflow.python.keras._impl.keras.models import Model 16 | 17 | from scale_layer_tf import Scale 18 | 19 | def identity_block(input_tensor, kernel_size, filters, stage, block): 20 | """ 21 | ### 说明: 22 | - The identity_block is the block that has no conv layer at shortcut 23 | 24 | ### Arguments: 25 | - input_tensor: input tensor 26 | - kernel_size: defualt 3, the kernel size of middle conv layer at main path 27 | - filters: list of integers, the nb_filters of 3 conv layer at main path 28 | - stage: integer, current stage label, used for generating layer names 29 | - block: 'a','b'..., current block label, used for generating layer names 30 | """ 31 | eps = 1.1e-5 32 | nb_filter1, nb_filter2, nb_filter3 = filters 33 | conv_name_base = 'res' + str(stage) + block + '_branch' 34 | bn_name_base = 'bn' + str(stage) + block + '_branch' 35 | scale_name_base = 'scale' + str(stage) + block + '_branch' 36 | 37 | tensor_x = Conv2D(nb_filter1, (1, 1), use_bias=False, name=conv_name_base + '2a')(input_tensor) 38 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name=bn_name_base + '2a')(tensor_x) 39 | tensor_x = Scale(axis=3, name=scale_name_base + '2a')(tensor_x) 40 | tensor_x = Activation('relu', name=conv_name_base + '2a_relu')(tensor_x) 41 | 42 | tensor_x = ZeroPadding2D((1, 1), name=conv_name_base + '2b_zeropadding')(tensor_x) 43 | tensor_x = Conv2D(nb_filter2, (kernel_size, kernel_size), use_bias=False, name=conv_name_base + '2b')(tensor_x) 44 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name=bn_name_base + '2b')(tensor_x) 45 | tensor_x = Scale(axis=3, name=scale_name_base + '2b')(tensor_x) 46 | tensor_x = Activation('relu', name=conv_name_base + '2b_relu')(tensor_x) 47 | 48 | tensor_x = Conv2D(nb_filter3, (1, 1), use_bias=False, name=conv_name_base + '2c')(tensor_x) 49 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name=bn_name_base + '2c')(tensor_x) 50 | tensor_x = Scale(axis=3, name=scale_name_base + '2c')(tensor_x) 51 | 52 | tensor_x = add([tensor_x, input_tensor], name='res' + str(stage) + block) 53 | tensor_x = Activation('relu', name='res' + str(stage) + block + '_relu')(tensor_x) 54 | return tensor_x 55 | 56 | def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)): 57 | ''' 58 | ### 说明 59 | - conv_block is the block that has a conv layer at shortcut 60 | 61 | ### Arguments 62 | - input_tensor: input tensor 63 | - kernel_size: defualt 3, the kernel size of middle conv layer at main path 64 | - filters: list of integers, the nb_filters of 3 conv layer at main path 65 | - stage: integer, current stage label, used for generating layer names 66 | - block: 'a','b'..., current block label, used for generating layer names 67 | 68 | ### 注意 69 | - Note that from stage 3, the first conv layer at main path is with strides=(2,2) 70 | - And the shortcut should have strides=(2,2) as well 71 | ''' 72 | eps = 1.1e-5 73 | nb_filter1, nb_filter2, nb_filter3 = filters 74 | conv_name_base = 'res' + str(stage) + block + '_branch' 75 | bn_name_base = 'bn' + str(stage) + block + '_branch' 76 | scale_name_base = 'scale' + str(stage) + block + '_branch' 77 | 78 | tensor_x = Conv2D(nb_filter1, (1, 1), strides=strides, use_bias=False, name=conv_name_base + '2a')(input_tensor) 79 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name=bn_name_base + '2a')(tensor_x) 80 | tensor_x = Scale(axis=3, name=scale_name_base + '2a')(tensor_x) 81 | tensor_x = Activation('relu', name=conv_name_base + '2a_relu')(tensor_x) 82 | 83 | tensor_x = ZeroPadding2D((1, 1), name=conv_name_base + '2b_zeropadding')(tensor_x) 84 | tensor_x = Conv2D(nb_filter2, (kernel_size, kernel_size), use_bias=False, name=conv_name_base + '2b')(tensor_x) 85 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name=bn_name_base + '2b')(tensor_x) 86 | tensor_x = Scale(axis=3, name=scale_name_base + '2b')(tensor_x) 87 | tensor_x = Activation('relu', name=conv_name_base + '2b_relu')(tensor_x) 88 | 89 | tensor_x = Conv2D(nb_filter3, (1, 1), use_bias=False, name=conv_name_base + '2c')(tensor_x) 90 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name=bn_name_base + '2c')(tensor_x) 91 | tensor_x = Scale(axis=3, name=scale_name_base + '2c')(tensor_x) 92 | 93 | shortcut = Conv2D(nb_filter3, (1, 1), strides=strides, use_bias=False, name=conv_name_base + '1')(input_tensor) 94 | shortcut = BatchNormalization(epsilon=eps, axis=3, name=bn_name_base + '1')(shortcut) 95 | shortcut = Scale(axis=3, name=scale_name_base + '1')(shortcut) 96 | 97 | tensor_x = add([tensor_x, shortcut], name='res' + str(stage) + block) 98 | tensor_x = Activation('relu', name='res' + str(stage) + block + '_relu')(tensor_x) 99 | 100 | return tensor_x 101 | 102 | def resnet_model(img_rows, img_cols, color_type=3, num_classes=3): 103 | ''' 104 | ### 说明: 105 | - DenseNet Model for Keras 106 | - Model Schema is based on https://github.com/flyyufelix/DenseNet-Keras 107 | - ImageNet Pretrained Weights 108 | 109 | ### Arguments: 110 | - img_rows: (int), image height 111 | - img_cols: (int), image width 112 | - color_type: (int), image channel 113 | - nb_layers: (list), number of filters in each block 114 | ### Returns 115 | - A Keras model instance. 116 | ''' 117 | 118 | eps = 1.1e-5 119 | 120 | # Handle Dimension Ordering for different backends 121 | img_input = Input(shape=(img_rows, img_cols, color_type), name='data') 122 | 123 | # conv1 124 | tensor_x = ZeroPadding2D((3, 3), name='conv1_zeropadding')(img_input) 125 | tensor_x = Conv2D(64, (7, 7), strides=(2, 2), use_bias=False, name='conv1')(tensor_x) 126 | tensor_x = BatchNormalization(epsilon=eps, axis=3, name='bn_conv1')(tensor_x) 127 | tensor_x = Scale(axis=3, name='scale_conv1')(tensor_x) 128 | tensor_x = Activation('relu', name='conv1_relu')(tensor_x) 129 | tensor_x = MaxPooling2D((3, 3), strides=(2, 2), name='pool1')(tensor_x) 130 | 131 | # conv2_x 132 | tensor_x = conv_block(tensor_x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1)) 133 | # for i in range(1, nb_layers[0]): 134 | # tensor_x = identity_block(tensor_x, 3, [64, 64, 256], stage=2, block='b' + str(i)) 135 | tensor_x = identity_block(tensor_x, 3, [64, 64, 256], stage=2, block='b') 136 | tensor_x = identity_block(tensor_x, 3, [64, 64, 256], stage=2, block='c') 137 | 138 | # conv3_x 139 | tensor_x = conv_block(tensor_x, 3, [128, 128, 512], stage=3, block='a') 140 | for i in range(1, 4): 141 | tensor_x = identity_block(tensor_x, 3, [128, 128, 512], stage=3, block='b' + str(i)) 142 | 143 | # conv4_x 144 | tensor_x = conv_block(tensor_x, 3, [256, 256, 1024], stage=4, block='a') 145 | for i in range(1, 23): 146 | tensor_x = identity_block(tensor_x, 3, [256, 256, 1024], stage=4, block='b' + str(i)) 147 | 148 | # conv5_x 149 | tensor_x = conv_block(tensor_x, 3, [512, 512, 2048], stage=5, block='a') 150 | # for i in range(1, nb_layers[3]): 151 | # tensor_x = identity_block(tensor_x, 3, [512, 512, 2048], stage=5, block='b' + str(i)) 152 | tensor_x = identity_block(tensor_x, 3, [512, 512, 2048], stage=5, block='b') 153 | tensor_x = identity_block(tensor_x, 3, [512, 512, 2048], stage=5, block='c') 154 | 155 | # fc 156 | x_fc = GlobalAveragePooling2D(name='avg_pool')(tensor_x) 157 | x_fc = Dense(num_classes, activation='softmax', name='fc3')(x_fc) 158 | 159 | model = Model(img_input, x_fc) 160 | 161 | # model.load_weights('resnet_101_finetune_weights_92%.h5') 162 | 163 | return model 164 | 165 | if __name__ == '__main__': 166 | 167 | model = resnet_model(480, 480, color_type=3, num_classes=3) 168 | model.summary() 169 | -------------------------------------------------------------------------------- /code/scale_layer_tf.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | # scale_layer_tf.py 5 | # @Author : yuanwenjin 6 | # @Mail : yfor1008@gmail.com 7 | # @Date : 2018-1-17 10:37:24 8 | # @Explanation : tensorflow keras实现 9 | """ 10 | 11 | from tensorflow.python.keras._impl.keras.layers.core import Layer 12 | # from tensorflow.python.keras._impl.keras.engine import Layer 13 | from tensorflow.python.keras._impl.keras.engine import InputSpec 14 | from tensorflow.python.keras._impl.keras import backend as K 15 | try: 16 | from tensorflow.python.keras._impl.keras import initializations 17 | except ImportError: 18 | from tensorflow.python.keras._impl.keras import initializers as initializations 19 | 20 | class Scale(Layer): 21 | '''Learns a set of weights and biases used for scaling the input data. 22 | the output consists simply in an element-wise multiplication of the input 23 | and a sum of a set of constants: 24 | 25 | out = in * gamma + beta, 26 | 27 | where 'gamma' and 'beta' are the weights and biases larned. 28 | 29 | # Arguments 30 | - axis: integer, axis along which to normalize in mode 0. For instance, 31 | if your input tensor has shape (samples, channels, rows, cols), 32 | set axis to 1 to normalize per feature map (channels axis). 33 | - momentum: momentum in the computation of the 34 | exponential average of the mean and standard deviation 35 | of the data, for feature-wise normalization. 36 | - weights: Initialization weights. 37 | List of 2 Numpy arrays, with shapes: 38 | `[(input_shape,), (input_shape,)]` 39 | - beta_init: name of initialization function for shift parameter 40 | (see [initializations](../initializations.md)), or alternatively, 41 | Theano/TensorFlow function to use for weights initialization. 42 | This parameter is only relevant if you don't pass a `weights` argument. 43 | - gamma_init: name of initialization function for scale parameter (see 44 | [initializations](../initializations.md)), or alternatively, 45 | Theano/TensorFlow function to use for weights initialization. 46 | This parameter is only relevant if you don't pass a `weights` argument. 47 | ''' 48 | def __init__(self, weights=None, axis=-1, momentum=0.9, beta_init='zero', gamma_init='one', **kwargs): 49 | self.momentum = momentum 50 | self.axis = axis 51 | self.beta_init = initializations.get(beta_init) 52 | self.gamma_init = initializations.get(gamma_init) 53 | self.beta = [] 54 | self.gamma = [] 55 | self.initial_weights = weights 56 | # self.trainable_weights = [] 57 | super(Scale, self).__init__(**kwargs) 58 | 59 | def build(self, input_shape): 60 | self.input_spec = [InputSpec(shape=input_shape)] 61 | shape = (int(input_shape[self.axis]),) 62 | 63 | # Compatibility with TensorFlow >= 1.0.0 64 | self.gamma = K.variable(self.gamma_init(shape), name='{}_gamma'.format(self.name)) 65 | self.beta = K.variable(self.beta_init(shape), name='{}_beta'.format(self.name)) 66 | #self.gamma = self.gamma_init(shape, name='{}_gamma'.format(self.name)) 67 | #self.beta = self.beta_init(shape, name='{}_beta'.format(self.name)) 68 | # self.trainable_weights = [self.gamma, self.beta] 69 | 70 | if self.initial_weights is not None: 71 | self.set_weights(self.initial_weights) 72 | del self.initial_weights 73 | 74 | def call(self, x): 75 | input_shape = self.input_spec[0].shape 76 | broadcast_shape = [1] * len(input_shape) 77 | broadcast_shape[self.axis] = input_shape[self.axis] 78 | 79 | out = K.reshape(self.gamma, broadcast_shape) * x + K.reshape(self.beta, broadcast_shape) 80 | return out 81 | 82 | def get_config(self): 83 | config = {"momentum": self.momentum, "axis": self.axis} 84 | base_config = super(Scale, self).get_config() 85 | return dict(list(base_config.items()) + list(config.items())) 86 | -------------------------------------------------------------------------------- /code/test_for_fp16.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | # test_for_fp16.py 5 | # @Author : yuanwenjin 6 | # @Mail : yfor1008@gmail.com 7 | # @Date : 2018-1-24 18:11:28 8 | # @Explanation : 测试tensorRT优化后效果 9 | """ 10 | 11 | import os 12 | import math 13 | import time 14 | import numpy as np 15 | import scipy.io as scio 16 | import cv2 17 | 18 | from tensorrt.lite import Engine 19 | import tensorrt as trt 20 | 21 | if __name__ == '__main__': 22 | 23 | # 模型 24 | # engine_fp16 = Engine(PLAN='./model/resnet_tf_b256_fp16.engine', log_sev=4, max_batch_size=256) 25 | # 屏蔽信息输出，log_sev应换成logger_severity 26 | engine_fp16 = Engine(PLAN='./model/resnet_tf_b256_fp16.engine', logger_severity=trt.infer.LogSeverity.ERROR, max_batch_size=256) 27 | 28 | 29 | # batch size 30 | batch_size = 256 31 | 32 | # 数据路径 33 | digestive_dir = u'/home/docker/code/app/tensorflow/digestive_change' 34 | 35 | data_dirs = [] 36 | for data_dir in os.listdir(digestive_dir): 37 | data_dirs.append(os.path.join(digestive_dir, data_dir)) 38 | 39 | for data_dir in data_dirs: 40 | imagelists = [img for img in os.listdir(data_dir) if img.endswith('jpg')] 41 | imagelists.sort() 42 | 43 | data_num = len(imagelists) 44 | # data_num = 10 45 | batch_num = int(math.ceil(float(data_num) / float(batch_size))) 46 | 47 | predictions = [] 48 | start_time = time.time() 49 | for batch_idx in range(batch_num): 50 | start = batch_idx * batch_size 51 | end = min((batch_idx + 1) * batch_size, data_num) 52 | 53 | # imgs = [] 54 | # for idx in range(start, end): 55 | # img = cv2.imread(os.path.join(data_dir, imagelists[idx]).encode('utf8')) 56 | # img = cv2.resize(img, (480, 480)) 57 | # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 58 | # img = img.transpose([2, 0, 1]) 59 | # imgs.append(img) 60 | imgs = [cv2.cvtColor(cv2.resize(cv2.imread(os.path.join(data_dir, imagelists[idx]).encode('utf8')), (480, 480)), 61 | cv2.COLOR_BGR2RGB).transpose([2, 0, 1]) for idx in range(start, end)] 62 | imgs = np.array(imgs, dtype='float32') / 255.0 63 | prediction = np.argmax(engine_fp16.infer(imgs)[0], 1).tolist() 64 | pre = [ff[0][0] for ff in prediction] 65 | predictions.extend(pre) 66 | 67 | end_time = time.time() 68 | with open(os.path.join('./results/', 'time_for_fp16.txt'), 'a') as ff: 69 | ff.write('%s: %f \n' % (os.path.basename(data_dir), end_time - start_time)) 70 | scio.savemat(os.path.join('./results/results_for_fp16', os.path.basename(data_dir) + '.mat'), {'data': predictions}) 71 | -------------------------------------------------------------------------------- /code/test_for_fp32.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | # test_for_RT.py 5 | # @Author : yuanwenjin 6 | # @Mail : yfor1008@gmail.com 7 | # @Date : 2018-1-24 15:45:44 8 | # @Explanation : 测试tensorRT优化后效果 9 | """ 10 | 11 | import os 12 | import math 13 | import time 14 | import numpy as np 15 | import scipy.io as scio 16 | import cv2 17 | 18 | from tensorrt.lite import Engine 19 | 20 | import tensorrt as trt 21 | 22 | if __name__ == '__main__': 23 | 24 | # 模型 25 | # engine_fp32 = Engine(PLAN='./model/resnet_tf_b256_fp32.engine', log_sev=4, max_batch_size=256) 26 | # 屏蔽信息输出，log_sev应换成logger_severity 27 | engine_fp32 = Engine(PLAN='./model/resnet_tf_b256_fp32.engine', logger_severity=trt.infer.LogSeverity.ERROR, max_batch_size=256) 28 | 29 | # batch size 30 | batch_size = 256 31 | 32 | # 数据路径 33 | digestive_dir = u'/home/docker/code/app/tensorflow/digestive_change' 34 | 35 | data_dirs = [] 36 | for data_dir in os.listdir(digestive_dir): 37 | data_dirs.append(os.path.join(digestive_dir, data_dir)) 38 | 39 | for data_dir in data_dirs: 40 | imagelists = [img for img in os.listdir(data_dir) if img.endswith('jpg')] 41 | imagelists.sort() 42 | 43 | data_num = len(imagelists) 44 | # data_num = 10 45 | batch_num = int(math.ceil(float(data_num) / float(batch_size))) 46 | 47 | predictions = [] 48 | start_time = time.time() 49 | for batch_idx in range(batch_num): 50 | start = batch_idx * batch_size 51 | end = min((batch_idx + 1) * batch_size, data_num) 52 | 53 | # imgs = [] 54 | # for idx in range(start, end): 55 | # img = cv2.imread(os.path.join(data_dir, imagelists[idx]).encode('utf8')) 56 | # img = cv2.resize(img, (480, 480)) 57 | # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 58 | # img = img.transpose([2, 0, 1]) 59 | # imgs.append(img) 60 | imgs = [cv2.cvtColor(cv2.resize(cv2.imread(os.path.join(data_dir, imagelists[idx]).encode('utf8')), (480, 480)), 61 | cv2.COLOR_BGR2RGB).transpose([2, 0, 1]) for idx in range(start, end)] 62 | imgs = np.array(imgs, dtype='float32') / 255.0 63 | prediction = np.argmax(engine_fp32.infer(imgs)[0], 1).tolist() 64 | pre = [ff[0][0] for ff in prediction] 65 | predictions.extend(pre) 66 | 67 | end_time = time.time() 68 | with open(os.path.join('./results/', 'time_for_fp32.txt'), 'a') as ff: 69 | ff.write('%s: %f \n' % (os.path.basename(data_dir), end_time - start_time)) 70 | scio.savemat(os.path.join('./results/results_for_fp32', os.path.basename(data_dir) + '.mat'), {'data': predictions}) 71 | -------------------------------------------------------------------------------- /code/test_for_keras.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | # test_for_keras.py 5 | # @Author : yuanwenjin 6 | # @Mail : yfor1008@gmail.com 7 | # @Date : 2018-1-24 15:04:51 8 | # @Explanation : 测试tensorflow中keras结果 9 | """ 10 | 11 | import os 12 | import math 13 | import time 14 | import numpy as np 15 | import scipy.io as scio 16 | import cv2 17 | from resnet_tf import resnet_model 18 | 19 | 20 | if __name__ == '__main__': 21 | 22 | # 生成模型 23 | model = resnet_model(480, 480, color_type=3, num_classes=3) 24 | # 导入训练好的模型 25 | model.load_weights('./model/resnet_101.h5') 26 | 27 | # batch size 28 | batch_size = 512 29 | 30 | # 数据路径 31 | digestive_dir = u'/home/docker/code/app/tensorflow/digestive_change' 32 | 33 | data_dirs = [] 34 | for data_dir in os.listdir(digestive_dir): 35 | data_dirs.append(os.path.join(digestive_dir, data_dir)) 36 | 37 | for data_dir in data_dirs: 38 | imagelists = [img for img in os.listdir(data_dir) if img.endswith('jpg')] 39 | imagelists.sort() 40 | 41 | data_num = len(imagelists) 42 | # data_num = 10 43 | batch_num = int(math.ceil(float(data_num) / float(batch_size))) 44 | 45 | predictions = [] 46 | start_time = time.time() 47 | for batch_idx in range(batch_num): 48 | start = batch_idx * batch_size 49 | end = min((batch_idx + 1) * batch_size, data_num) 50 | 51 | imgs = [] 52 | for idx in range(start, end): 53 | img = cv2.imread(os.path.join(data_dir, imagelists[idx]).encode('utf8')) 54 | img = cv2.resize(img, (480, 480)) 55 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 56 | imgs.append(img) 57 | imgs = np.array(imgs, dtype='float32') / 255.0 58 | prediction = np.argmax(model.predict(np.array(imgs)), 1).tolist() 59 | predictions.extend(prediction) 60 | 61 | end_time = time.time() 62 | with open(os.path.join('./results/', 'time_for_keras.txt'), 'a') as ff: 63 | ff.write('%s: %f \n' % (os.path.basename(data_dir), end_time - start_time)) 64 | scio.savemat(os.path.join('./results/results_for_keras', os.path.basename(data_dir) + '.mat'), {'data': predictions}) 65 | -------------------------------------------------------------------------------- /code/train.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | # train.py 5 | # @Author : yuanwenjin 6 | # @Mail : yfor1008@gmail.com 7 | # @Date : 2018-1-17 11:11:06 8 | # @Explanation : 说明 9 | """ 10 | 11 | import tensorflow as tf 12 | from resnet_tf import resnet_model 13 | 14 | from tensorflow.python.keras._impl.keras.optimizers import SGD 15 | from tensorflow.python.keras._impl.keras.preprocessing.image import ImageDataGenerator 16 | from tensorflow.python.keras._impl.keras import backend as K 17 | from tensorflow.python.keras._impl.keras.callbacks import ModelCheckpoint 18 | from tensorflow.python.keras._impl.keras.callbacks import EarlyStopping 19 | from tensorflow.python.keras._impl.keras.callbacks import LearningRateScheduler 20 | from tensorflow.python.keras._impl.keras.callbacks import TensorBoard 21 | 22 | from tensorflow.python.framework import graph_io 23 | from tensorflow.python.tools import freeze_graph 24 | from tensorflow.core.protobuf import saver_pb2 25 | from tensorflow.python.training import saver as saver_lib 26 | 27 | if __name__ == '__main__': 28 | 29 | config = { 30 | # Training params 31 | "train_data_dir": "/home/docker/code/app/get_samples/samples_train/train", # training data 32 | "val_data_dir": "/home/docker/code/app/get_samples/samples_train/valid", # validation data 33 | "train_batch_size": 8, # training batch size 34 | "epochs": 30, # number of training epochs 35 | "num_train_samples" : 217728, # number of training examples 36 | "num_val_samples" : 55872, # number of test examples 37 | # "num_train_samples" : 500, # number of training examples 38 | # "num_val_samples" : 300, # number of test examples 39 | 40 | # Where to save models (Tensorflow + TensorRT) 41 | "graphdef_file": "./model/resnet_tf_graphdef.pb", 42 | "frozen_model_file": "./model/resnet_tf_frozen_model.pb", 43 | "snapshot_dir": "./model/snapshot", 44 | "engine_save_dir": "./model/", 45 | 46 | # Needed for TensorRT 47 | "image_dim": 480, # the image size (square images) 48 | "inference_batch_size": 1, # inference batch size 49 | "input_layer": "data", # name of the input tensor in the TF computational graph 50 | "out_layer": "fc3/Softmax", # name of the output tensorf in the TF conputational graph 51 | "output_size" : 3, # number of classes in output (3) 52 | "precision": "fp32", # desired precision (fp32, fp16) 53 | 54 | "test_image_path" : "/home/data/val/roses" 55 | } 56 | 57 | model = resnet_model(480, 480, color_type=3, num_classes=3) 58 | model.summary() 59 | 60 | # 导入预先训练好的模型，进行finetune 61 | model.load_weights('./model/resnet_101.h5') 62 | 63 | # 数据保存 64 | cbks = [TensorBoard(log_dir='./logs', histogram_freq=0, write_images=True, write_graph=True, write_grads=True), 65 | ModelCheckpoint('./model/resnet_101_weights.h5', monitor='val_loss', verbose=1, save_weights_only=True), 66 | ] 67 | 68 | # 优化参数 69 | sgd = SGD(lr=1e-4, decay=1e-6, momentum=0.9, nesterov=True) 70 | model.compile(optimizer=sgd, loss='categorical_crossentropy', metrics=['accuracy']) 71 | 72 | data_gen = ImageDataGenerator(rescale=1./255, 73 | shear_range=10, 74 | width_shift_range=0.1, 75 | height_shift_range=0.1) 76 | train_generator = data_gen.flow_from_directory(directory=config['train_data_dir'], 77 | target_size=(config['image_dim'], config['image_dim']), 78 | batch_size=config['train_batch_size'], 79 | class_mode='categorical') 80 | val_generator = data_gen.flow_from_directory(directory=config['val_data_dir'], 81 | target_size=(config['image_dim'], config['image_dim']), 82 | batch_size=config['train_batch_size'], 83 | class_mode='categorical') 84 | 85 | # train the model on the new data for a few epochs 86 | model.fit_generator(train_generator, 87 | steps_per_epoch=config['num_train_samples']//config['train_batch_size'], 88 | epochs=config['epochs'], 89 | # verbose=1, 90 | # callbacks=cbks, 91 | validation_data=val_generator, 92 | validation_steps=config['num_val_samples']//config['train_batch_size'], 93 | callbacks=cbks, 94 | verbose=1) 95 | 96 | model.save_weights('./model/resnet_101_last.h5') 97 | 98 | # Now, let's use the Tensorflow backend to get the TF graphdef and frozen graph 99 | K.set_learning_phase(0) 100 | sess = K.get_session() 101 | saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2) 102 | 103 | # save model weights in TF checkpoint 104 | checkpoint_path = saver.save(sess, config['snapshot_dir'], global_step=0, latest_filename='checkpoint_state') 105 | # checkpoint_path = saver.restore(sess, config['snapshot_dir']) 106 | 107 | # remove nodes not needed for inference from graph def 108 | train_graph = sess.graph 109 | inference_graph = tf.graph_util.remove_training_nodes(train_graph.as_graph_def()) 110 | 111 | # write the graph definition to a file. 112 | # You can view this file to see your network structure and 113 | # to determine the names of your network's input/output layers. 114 | graph_io.write_graph(inference_graph, '.', config['graphdef_file']) 115 | 116 | # specify which layer is the output layer for your graph. 117 | # In this case, we want to specify the softmax layer after our 118 | # last dense (fully connected) layer. 119 | out_names = config['out_layer'] 120 | 121 | # freeze your inference graph and save it for later! (Tensorflow) 122 | freeze_graph.freeze_graph(config['graphdef_file'], 123 | '', 124 | False, 125 | checkpoint_path, 126 | out_names, 127 | "save/restore_all", 128 | "save/Const:0", 129 | config['frozen_model_file'], 130 | False, 131 | "" 132 | ) 133 | -------------------------------------------------------------------------------- /code/train_densenet_tf.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | # train_densenet_tf.py 5 | # @Author : yuanwenjin 6 | # @Mail : yfor1008@gmail.com 7 | # @Date : 2018-3-1 17:08:14 8 | # @Explanation : 训练densenet_tf 9 | """ 10 | 11 | import tensorflow as tf 12 | from densenet_tf import generate_densenet 13 | 14 | from tensorflow.python.keras._impl.keras.optimizers import SGD 15 | from tensorflow.python.keras._impl.keras.preprocessing.image import ImageDataGenerator 16 | from tensorflow.python.keras._impl.keras import backend as K 17 | from tensorflow.python.keras._impl.keras.callbacks import ModelCheckpoint 18 | from tensorflow.python.keras._impl.keras.callbacks import EarlyStopping 19 | from tensorflow.python.keras._impl.keras.callbacks import LearningRateScheduler 20 | from tensorflow.python.keras._impl.keras.callbacks import TensorBoard 21 | 22 | from tensorflow.python.framework import graph_io 23 | from tensorflow.python.tools import freeze_graph 24 | from tensorflow.core.protobuf import saver_pb2 25 | from tensorflow.python.training import saver as saver_lib 26 | 27 | if __name__ == '__main__': 28 | 29 | model_size = 121 30 | num_classes = 3 31 | 32 | config = { 33 | # Training params 34 | "train_data_dir": "/home/docker/code/app/get_samples/samples_train/train", # training data 35 | "val_data_dir": "/home/docker/code/app/get_samples/samples_train/valid", # validation data 36 | "train_batch_size": 8, # training batch size 37 | "epochs": 3, # number of training epochs 38 | # "num_train_samples" : 217728, # number of training examples 39 | # "num_val_samples" : 55872, # number of test examples 40 | "num_train_samples" : 500, # number of training examples 41 | "num_val_samples" : 300, # number of test examples 42 | 43 | # Where to save models (Tensorflow + TensorRT) 44 | "graphdef_file": "./model_densenet/densenet_tf_graphdef.pb", 45 | "frozen_model_file": "./model_densenet/densenet_tf_frozen_model.pb", 46 | "snapshot_dir": "./model_densenet/snapshot", 47 | "engine_save_dir": "./model_densenet/", 48 | 49 | # Needed for TensorRT 50 | "image_dim": 480, # the image size (square images) 51 | "inference_batch_size": 1, # inference batch size 52 | "input_layer": "data", # name of the input tensor in the TF computational graph 53 | "out_layer": "fc6/Softmax", # name of the output tensorf in the TF conputational graph 54 | "output_size" : num_classes, # number of classes in output (3) 55 | "precision": "fp32", # desired precision (fp32, fp16) 56 | 57 | } 58 | 59 | model = generate_densenet(480, 480, color_type=3, model_size=model_size, num_classes=num_classes) 60 | model.summary() 61 | 62 | # # 导入预先训练好的模型，进行finetune 63 | # model.load_weights('./model_densenet/densenet_%s_weights.h5' % (model_size)) 64 | 65 | # 数据保存 66 | cbks = [TensorBoard(log_dir='./logs_densenet', histogram_freq=0, write_images=True, write_graph=True, write_grads=True), 67 | ModelCheckpoint('./model_densenet/resnet_%s_weights.h5' % model_size, monitor='val_loss', verbose=1, save_weights_only=True), 68 | ] 69 | 70 | # 优化参数 71 | sgd = SGD(lr=1e-4, decay=1e-4, momentum=0.9, nesterov=True) 72 | model.compile(optimizer=sgd, loss='categorical_crossentropy', metrics=['accuracy']) 73 | 74 | data_gen = ImageDataGenerator(rescale=1./255, 75 | shear_range=10, 76 | width_shift_range=0.1, 77 | height_shift_range=0.1) 78 | train_generator = data_gen.flow_from_directory(directory=config['train_data_dir'], 79 | target_size=(config['image_dim'], config['image_dim']), 80 | batch_size=config['train_batch_size'], 81 | class_mode='categorical') 82 | val_generator = data_gen.flow_from_directory(directory=config['val_data_dir'], 83 | target_size=(config['image_dim'], config['image_dim']), 84 | batch_size=config['train_batch_size'], 85 | class_mode='categorical') 86 | 87 | # train the model on the new data for a few epochs 88 | model.fit_generator(train_generator, 89 | steps_per_epoch=config['num_train_samples']//config['train_batch_size'], 90 | epochs=config['epochs'], 91 | # verbose=1, 92 | # callbacks=cbks, 93 | validation_data=val_generator, 94 | validation_steps=config['num_val_samples']//config['train_batch_size'], 95 | callbacks=cbks, 96 | verbose=1) 97 | 98 | model.save_weights('./model_densenet/resnet_%s_last.h5' % model_size) 99 | 100 | # Now, let's use the Tensorflow backend to get the TF graphdef and frozen graph 101 | K.set_learning_phase(0) 102 | sess = K.get_session() 103 | saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2) 104 | 105 | # save model weights in TF checkpoint 106 | checkpoint_path = saver.save(sess, config['snapshot_dir'], global_step=0, latest_filename='checkpoint_state') 107 | # checkpoint_path = saver.restore(sess, config['snapshot_dir']) 108 | 109 | # remove nodes not needed for inference from graph def 110 | train_graph = sess.graph 111 | inference_graph = tf.graph_util.remove_training_nodes(train_graph.as_graph_def()) 112 | 113 | # write the graph definition to a file. 114 | # You can view this file to see your network structure and 115 | # to determine the names of your network's input/output layers. 116 | graph_io.write_graph(inference_graph, '.', config['graphdef_file']) 117 | 118 | # specify which layer is the output layer for your graph. 119 | # In this case, we want to specify the softmax layer after our 120 | # last dense (fully connected) layer. 121 | out_names = config['out_layer'] 122 | 123 | # freeze your inference graph and save it for later! (Tensorflow) 124 | freeze_graph.freeze_graph(config['graphdef_file'], 125 | '', 126 | False, 127 | checkpoint_path, 128 | out_names, 129 | "save/restore_all", 130 | "save/Const:0", 131 | config['frozen_model_file'], 132 | False, 133 | "" 134 | ) 135 | --------------------------------------------------------------------------------