30
30
from object_detection .models import faster_rcnn_inception_resnet_v2_feature_extractor as frcnn_inc_res
31
31
from object_detection .models import faster_rcnn_inception_v2_feature_extractor as frcnn_inc_v2
32
32
from object_detection .models import faster_rcnn_nas_feature_extractor as frcnn_nas
33
+ from object_detection .models import faster_rcnn_pnas_feature_extractor as frcnn_pnas
33
34
from object_detection .models import faster_rcnn_resnet_v1_feature_extractor as frcnn_resnet_v1
34
35
from object_detection .models import ssd_resnet_v1_fpn_feature_extractor as ssd_resnet_v1_fpn
35
36
from object_detection .models .embedded_ssd_mobilenet_v1_feature_extractor import EmbeddedSSDMobileNetV1FeatureExtractor
55
56
FASTER_RCNN_FEATURE_EXTRACTOR_CLASS_MAP = {
56
57
'faster_rcnn_nas' :
57
58
frcnn_nas .FasterRCNNNASFeatureExtractor ,
59
+ 'faster_rcnn_pnas' :
60
+ frcnn_pnas .FasterRCNNPNASFeatureExtractor ,
58
61
'faster_rcnn_inception_resnet_v2' :
59
62
frcnn_inc_res .FasterRCNNInceptionResnetV2FeatureExtractor ,
60
63
'faster_rcnn_inception_v2' :
@@ -95,13 +98,19 @@ def build(model_config, is_training, add_summaries=True):
95
98
96
99
97
100
def _build_ssd_feature_extractor (feature_extractor_config , is_training ,
98
- reuse_weights = None ):
101
+ reuse_weights = None ,
102
+ inplace_batchnorm_update = False ):
99
103
"""Builds a ssd_meta_arch.SSDFeatureExtractor based on config.
100
104
101
105
Args:
102
106
feature_extractor_config: A SSDFeatureExtractor proto config from ssd.proto.
103
107
is_training: True if this feature extractor is being built for training.
104
108
reuse_weights: if the feature extractor should reuse weights.
109
+ inplace_batchnorm_update: Whether to update batch_norm inplace during
110
+ training. This is required for batch norm to work correctly on TPUs. When
111
+ this is false, user must add a control dependency on
112
+ tf.GraphKeys.UPDATE_OPS for train/loss op in order to update the batch
113
+ norm moving average parameters.
105
114
106
115
Returns:
107
116
ssd_meta_arch.SSDFeatureExtractor based on config.
@@ -126,7 +135,8 @@ def _build_ssd_feature_extractor(feature_extractor_config, is_training,
126
135
return feature_extractor_class (is_training , depth_multiplier , min_depth ,
127
136
pad_to_multiple , conv_hyperparams ,
128
137
batch_norm_trainable , reuse_weights ,
129
- use_explicit_padding , use_depthwise )
138
+ use_explicit_padding , use_depthwise ,
139
+ inplace_batchnorm_update )
130
140
131
141
132
142
def _build_ssd_model (ssd_config , is_training , add_summaries ):
@@ -140,15 +150,18 @@ def _build_ssd_model(ssd_config, is_training, add_summaries):
140
150
141
151
Returns:
142
152
SSDMetaArch based on the config.
153
+
143
154
Raises:
144
155
ValueError: If ssd_config.type is not recognized (i.e. not registered in
145
156
model_class_map).
146
157
"""
147
158
num_classes = ssd_config .num_classes
148
159
149
160
# Feature extractor
150
- feature_extractor = _build_ssd_feature_extractor (ssd_config .feature_extractor ,
151
- is_training )
161
+ feature_extractor = _build_ssd_feature_extractor (
162
+ feature_extractor_config = ssd_config .feature_extractor ,
163
+ is_training = is_training ,
164
+ inplace_batchnorm_update = ssd_config .inplace_batchnorm_update )
152
165
153
166
box_coder = box_coder_builder .build (ssd_config .box_coder )
154
167
matcher = matcher_builder .build (ssd_config .matcher )
@@ -194,21 +207,29 @@ def _build_ssd_model(ssd_config, is_training, add_summaries):
194
207
195
208
196
209
def _build_faster_rcnn_feature_extractor (
197
- feature_extractor_config , is_training , reuse_weights = None ):
210
+ feature_extractor_config , is_training , reuse_weights = None ,
211
+ inplace_batchnorm_update = False ):
198
212
"""Builds a faster_rcnn_meta_arch.FasterRCNNFeatureExtractor based on config.
199
213
200
214
Args:
201
215
feature_extractor_config: A FasterRcnnFeatureExtractor proto config from
202
216
faster_rcnn.proto.
203
217
is_training: True if this feature extractor is being built for training.
204
218
reuse_weights: if the feature extractor should reuse weights.
219
+ inplace_batchnorm_update: Whether to update batch_norm inplace during
220
+ training. This is required for batch norm to work correctly on TPUs. When
221
+ this is false, user must add a control dependency on
222
+ tf.GraphKeys.UPDATE_OPS for train/loss op in order to update the batch
223
+ norm moving average parameters.
205
224
206
225
Returns:
207
226
faster_rcnn_meta_arch.FasterRCNNFeatureExtractor based on config.
208
227
209
228
Raises:
210
229
ValueError: On invalid feature extractor type.
211
230
"""
231
+ if inplace_batchnorm_update :
232
+ raise ValueError ('inplace batchnorm updates not supported.' )
212
233
feature_type = feature_extractor_config .type
213
234
first_stage_features_stride = (
214
235
feature_extractor_config .first_stage_features_stride )
@@ -238,6 +259,7 @@ def _build_faster_rcnn_model(frcnn_config, is_training, add_summaries):
238
259
239
260
Returns:
240
261
FasterRCNNMetaArch based on the config.
262
+
241
263
Raises:
242
264
ValueError: If frcnn_config.type is not recognized (i.e. not registered in
243
265
model_class_map).
@@ -246,7 +268,8 @@ def _build_faster_rcnn_model(frcnn_config, is_training, add_summaries):
246
268
image_resizer_fn = image_resizer_builder .build (frcnn_config .image_resizer )
247
269
248
270
feature_extractor = _build_faster_rcnn_feature_extractor (
249
- frcnn_config .feature_extractor , is_training )
271
+ frcnn_config .feature_extractor , is_training ,
272
+ frcnn_config .inplace_batchnorm_update )
250
273
251
274
number_of_stages = frcnn_config .number_of_stages
252
275
first_stage_anchor_generator = anchor_generator_builder .build (
0 commit comments