def get_models():
model_cropped = resnet.resnet_v1(resnet_depth=FLAGS.resnet_depth,
width_multiplier=FLAGS.width_multiplier,
cifar_stem=FLAGS.crop_size <= 32)
# generic resnet_v2 keras
model_full = resnet.resnet_v1(resnet_depth=FLAGS.resnet_depth,
width_multiplier=FLAGS.width_multiplier,
cifar_stem=FLAGS.image_size <= 32)
# # with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
# # uses call(inputs, training) instead of call(inputs, is_training)
# # does call() in keras use call(imputs, is_training)? (not tf.keras, just keras) No, they are now in sync
# # tf.disable_eager_execution()
# # with tf.Graph().as_default():
# model_cropped = tf2.keras.applications.resnet.ResNet50(include_top=False, weights=None, input_tensor=None,
# input_shape=(FLAGS.image_size, FLAGS.image_size, 3),
# pooling=None)
# model_full = tf2.keras.applications.resnet.ResNet50(include_top=False, weights=None, input_tensor=None,
# input_shape=(FLAGS.image_size, FLAGS.image_size, 3),
# pooling=None)
# model_cropped.summary()
# model_full.summary()
# # ValueError: Tensor("conv1_conv_1/kernel/Read/ReadVariableOp:0", shape=(7, 7, 3, 64), dtype=float32) must be from the same graph as Tensor("base_model/resnet50/conv1_pad/Pad:0", shape=(128, 38, 38, 3), dtype=float32).
models = {'model_full': model_full, 'model_cropped': model_cropped}
return models
def resnet_attention_train_predict(row, col, modelfile, info):
(x_train, y_train), (x_test, y_test), (x_indt, y_indt) = load_data(col=col, row=row)
model = resnet_v1(input_shape=(row, col+1, 2), depth=20, num_classes=2)
model.compile(optimizer=Adam(learning_rate=lr_schedule(0)),
loss='categorical_crossentropy',
metrics=['accuracy'])
lr_scheduler = LearningRateScheduler(lr_schedule)
lr_reducer = ReduceLROnPlateau(factor=np.sqrt(0.1), cooldown=0, patience=5,
min_lr=0.5e-6)
'''checkpoint = ModelCheckpoint(filepath=modelfile, monitor='val_accuracy',
verbose=1, save_best_only=True,
save_weights_only=True)
cbs = [checkpoint, lr_reducer, lr_scheduler]
model.fit(x_train, y_train,
batch_size=32,
epochs=10,
validation_data=[x_test, y_test],
shuffle=True,
callbacks=cbs)'''
model.load_weights(modelfile)
y_pred = model.predict(x_test)
info += 'By resnet with attention Training in TRB, Testing in testdata.....'
writeMetrics('TRB_result.txt', y_test, y_pred, info)
indt_pred = model.predict(x_indt)
info += 'By resnet with attention Tringing in TRB, Testing in nCoV-19...'
writeMetrics('TRB_result.txt', y_indt, indt_pred, info)
def resnet_train_predict():
(x_train, y_train), (x_test, y_test) = load_data()
model = resnet_v1(input_shape=(30,16), depth=20, num_classes=2)
model.compile(loss='categorical_crossentropy',
optimizer=Adam(learning_rate=lr_schedule(0)),
metrics=['accuracy'])
lr_scheduler = LearningRateScheduler(lr_schedule)
lr_reducer = ReduceLROnPlateau(factor=np.sqrt(0.1), cooldown=0, patience=5,
min_lr=0.5e-6)
save_dir = os.path.join(os.getcwd(), 'save_models')
model_name = 'trb_resnet_model.{epoch:02d}.h5'
if not os.path.isdir(save_dir):
os.makedirs(save_dir)
filepath = os.path.join(save_dir, model_name)
checkpoint = ModelCheckpoint(filepath=filepath, monitor='val_accuracy',
verbose=1, save_best_only=True,
save_weights_only=True)
cbs = [checkpoint, lr_reducer, lr_scheduler]
model.fit(x_train, y_train, batch_size=32, epochs=20,
validation_data=[x_test, y_test],
shuffle=True,
callbacks=cbs)
y_pred = model.predict(x_test, batch_size=48)
info = 'Training in TRB.....'
writeMetrics('TRB_result.txt', y_test, y_pred, info)
def resnet_fpn(
features,
min_level=3,
max_level=7, # pylint: disable=unused-argument
resnet_depth=50,
conv0_kernel_size=7,
conv0_space_to_depth_block_size=0,
is_training_bn=False):
"""ResNet feature pyramid networks."""
# upward layers
with tf.variable_scope('resnet%s' % resnet_depth):
resnet_fn = resnet.resnet_v1(resnet_depth, conv0_kernel_size,
conv0_space_to_depth_block_size)
u2, u3, u4, u5 = resnet_fn(features, is_training_bn)
feats_bottom_up = {
2: u2,
3: u3,
4: u4,
5: u5,
}
with tf.variable_scope('resnet_fpn'):
# lateral connections
feats_lateral = {}
for level in range(min_level, _RESNET_MAX_LEVEL + 1):
feats_lateral[level] = tf.layers.conv2d(feats_bottom_up[level],
filters=256,
kernel_size=(1, 1),
padding='same',
name='l%d' % level)
# add top-down path
feats = {_RESNET_MAX_LEVEL: feats_lateral[_RESNET_MAX_LEVEL]}
for level in range(_RESNET_MAX_LEVEL - 1, min_level - 1, -1):
feats[level] = nearest_upsampling(feats[level + 1],
2) + feats_lateral[level]
# add post-hoc 3x3 convolution kernel
for level in range(min_level, _RESNET_MAX_LEVEL + 1):
feats[level] = tf.layers.conv2d(feats[level],
filters=256,
strides=(1, 1),
kernel_size=(3, 3),
padding='same',
name='post_hoc_d%d' % level)
# Use original FPN P6 level implementation from CVPR'17 FPN paper instead of
# coarse FPN levels introduced for RetinaNet.
# Reference: https://github.com/ddkang/Detectron/blob/80f329530843e66d07ca39e19901d5f3e5daf009/lib/modeling/FPN.py # pylint: disable=line-too-long
feats[6] = tf.layers.max_pooling2d(inputs=feats[5],
pool_size=1,
strides=2,
padding='valid',
name='p6')
return feats
def load_model(self, model_id=0, weights_file='cifar10.hdf5'):
model = resnet_v1(input_shape=self.input_shape, depth=self.depth)
model.compile(loss='categorical_crossentropy',
optimizer=Adam(lr=lr_schedule(self.epoch)),
metrics=['accuracy'])
weights_file = os.path.join(self.script_path, weights_file)
if not os.path.isfile(weights_file):
logger.fatal("You have not train the model yet, please train model first.")
model.load_weights(weights_file)
return model
def train_dnn_model(self, _model=None, x_train=None, y_train=None,
x_val=None, y_val=None, train_strategy=None):
"""train a dnn model on cifar-10 dataset based on train strategy"""
# k.set_image_data_format('channels_last')
# model_id = _model[0]
weights_file = _model[1]
weights_file = os.path.join(self.script_path, weights_file)
model = resnet_v1(input_shape=self.input_shape, depth=self.depth)
def categorical_crossentropy_wrapper(y_true, y_pred):
y_pred = keras.backend.clip(y_pred, self.min_value, self.max_value)
return keras.losses.categorical_crossentropy(y_true, y_pred)
model.compile(loss=categorical_crossentropy_wrapper,
optimizer=Adam(lr=lr_schedule(0)),
metrics=['accuracy'])
lr_scheduler = LearningRateScheduler(lr_schedule)
lr_reducer = ReduceLROnPlateau(factor=np.sqrt(0.1),
cooldown=0,
patience=5,
min_lr=0.5e-6)
checkpoint = ModelCheckpoint(weights_file, monitor='acc', verbose=1,
save_best_only=False)
callbacks_list = [checkpoint, lr_reducer, lr_scheduler]
x_val = self.preprocess_original_imgs(x_val)
if train_strategy is None:
x_train = self.preprocess_original_imgs(x_train)
self.datagen_rotation.fit(x_train)
data = self.datagen_rotation.flow(x_train, y_train, batch_size=self.batch_size)
# model.fit(x_train, y_train,
# batch_size=self.batch_size,
# epochs=self.epoch,
# validation_data=(x_val, y_val),
# shuffle=True,
# callbacks=callbacks_list)
else:
graph = tf.get_default_graph()
data = DataGenerator(self, model, x_train, y_train, self.batch_size, train_strategy, graph)
model.fit_generator(data,
validation_data=(x_val, y_val),
epochs=self.epoch, verbose=1, workers=4,
callbacks=callbacks_list)
return model
def main():
# TODO record the learning rate vs epoch
cifar10 = keras.datasets.cifar10
(train_images, train_labels), (test_images, test_labels) = cifar10.load_data()
train_labels = keras.utils.to_categorical(train_labels)
input_shape = train_images.shape[1:]
batch_size = 32
num_classes = 10
epochs = 100
data_augmentation = True
num_predictions = 20
# model type
selected_model = "ResNet20v2"
# selected_model = "keras.applications.ResNet50V2"
n = 2 # order of ResNetv2, 2 or 6
version = 2
depth = model_depth(n, version)
model_type = 'ResNet%dv%d' % (depth, version)
save_dir = os.path.join(os.getcwd(), 'saved_models')
model_name = 'keras_cifar10_trained_model.h5'
if version == 2:
model = resnet_v2(input_shape=input_shape, depth=depth)
else:
model = resnet_v1(input_shape=input_shape, depth=depth)
model.compile(loss='categorical_crossentropy',
optimizer=Adam(lr=lr_schedule(0)),
metrics=['accuracy'])
# callbacks
logdir = os.path.join(
"logs", datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
csv_logger = CSVLogger(os.path.join(
logdir, "training.log.csv"), append=True)
tensorboard_callback = tf.keras.callbacks.TensorBoard(
logdir, histogram_freq=1)
callbacks = [csv_logger, tensorboard_callback]
model.fit(train_images,
train_labels,
epochs=epochs,
validation_data=(test_images, test_labels),
batch_size=batch_size,
verbose=1, workers=4,
callbacks=callbacks)
def create_module_from_checkpoints(args):
resnet.BATCH_NORM_DECAY = FLAGS.batch_norm_decay
model = resnet.resnet_v1(
resnet_depth=FLAGS.resnet_depth,
width_multiplier=FLAGS.width_multiplier,
cifar_stem=FLAGS.image_size <= 32)
#save the model in the same folder as the checkpoints
print('Start: Creating Hub Module from FLAGS.checkpoint_path')
#global_step = int(FLAGS.checkpoint_path[-1])
hub_name = os.path.basename(FLAGS.checkpoint_path) # Global step is actually the hub folder name used
hub_export_dir = build_hub_module(model, FLAGS.num_classes,
global_step = hub_name,
checkpoint_path= FLAGS.checkpoint_path)
print('Hub Module Created')
def model_fn_new(img_shape, num_classes, learning_rate):
# Subtracting pixel mean improves accuracy
# Model parameter
# ----------------------------------------------------------------------------
# | | 200-epoch | Orig Paper| 200-epoch | Orig Paper| sec/epoch
# Model | n | ResNet v1 | ResNet v1 | ResNet v2 | ResNet v2 | GTX1080Ti
# |v1(v2)| %Accuracy | %Accuracy | %Accuracy | %Accuracy | v1 (v2)
# ----------------------------------------------------------------------------
# ResNet20 | 3 (2)| 92.16 | 91.25 | ----- | ----- | 35 (---)
# ResNet32 | 5(NA)| 92.46 | 92.49 | NA | NA | 50 ( NA)
# ResNet44 | 7(NA)| 92.50 | 92.83 | NA | NA | 70 ( NA)
# ResNet56 | 9 (6)| 92.71 | 93.03 | 93.01 | NA | 90 (100)
# ResNet110 |18(12)| 92.65 | 93.39+-.16| 93.15 | 93.63 | 165(180)
# ResNet164 |27(18)| ----- | 94.07 | ----- | 94.54 | ---(---)
# ResNet1001| (111)| ----- | 92.39 | ----- | 95.08+-.14| ---(---)
# ---------------------------------------------------------------------------
n = 3
# Model version
# Orig paper: version = 1 (ResNet v1), Improved ResNet: version = 2 (ResNet v2)
version = 1
# Computed depth from supplied model parameter n
if version == 1:
depth = n * 6 + 2
elif version == 2:
depth = n * 9 + 2
from resnet import resnet_v1, resnet_v2
if version == 2:
model = resnet_v2(input_shape=img_shape, depth=depth)
else:
model = resnet_v1(input_shape=img_shape, depth=depth)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
model.compile(loss='categorical_crossentropy',
optimizer=optimizer,
metrics=metrics)
info('ResNet%dv%d' % (depth, version))
return model
def _model_outputs():
"""Generates outputs from the model."""
model_outputs = {}
with tf.variable_scope('resnet%s' % params['resnet_depth']):
resnet_fn = resnet.resnet_v1(
params['resnet_depth'],
num_batch_norm_group=params['num_batch_norm_group'])
backbone_feats = resnet_fn(features['images'],
params['is_training_bn'])
fpn_feats = fpn.fpn(backbone_feats, params['min_level'],
params['max_level'])
rpn_score_outputs, rpn_box_outputs = heads.rpn_head(
fpn_feats, params['min_level'], params['max_level'],
len(params['aspect_ratios'] * params['num_scales']))
if mode == tf.estimator.ModeKeys.TRAIN:
rpn_pre_nms_topn = params['rpn_pre_nms_topn']
rpn_post_nms_topn = params['rpn_post_nms_topn']
else:
rpn_pre_nms_topn = params['test_rpn_pre_nms_topn']
rpn_post_nms_topn = params['test_rpn_post_nms_topn']
_, rpn_box_rois = mask_rcnn_architecture.proposal_op(
rpn_score_outputs, rpn_box_outputs, all_anchors,
features['image_info'], rpn_pre_nms_topn, rpn_post_nms_topn,
params['rpn_nms_threshold'], params['rpn_min_size'])
rpn_box_rois = tf.to_float(rpn_box_rois)
if mode == tf.estimator.ModeKeys.TRAIN:
# Sampling
box_targets, class_targets, rpn_box_rois, proposal_to_label_map = (
mask_rcnn_architecture.proposal_label_op(
rpn_box_rois,
labels['gt_boxes'],
labels['gt_classes'],
features['image_info'],
batch_size_per_im=params['batch_size_per_im'],
fg_fraction=params['fg_fraction'],
fg_thresh=params['fg_thresh'],
bg_thresh_hi=params['bg_thresh_hi'],
bg_thresh_lo=params['bg_thresh_lo']))
# Performs multi-level RoIAlign.
box_roi_features = ops.multilevel_crop_and_resize(fpn_feats,
rpn_box_rois,
output_size=7)
class_outputs, box_outputs = heads.box_head(
box_roi_features,
num_classes=params['num_classes'],
mlp_head_dim=params['fast_rcnn_mlp_head_dim'])
if mode != tf.estimator.ModeKeys.TRAIN:
batch_size, _, _ = class_outputs.get_shape().as_list()
detections = []
softmax_class_outputs = tf.nn.softmax(class_outputs)
for i in range(batch_size):
detections.append(
anchors.generate_detections_per_image_op(
softmax_class_outputs[i], box_outputs[i],
rpn_box_rois[i], features['source_ids'][i],
features['image_info'][i],
params['test_detections_per_image'],
params['test_rpn_post_nms_topn'], params['test_nms'],
params['bbox_reg_weights']))
detections = tf.stack(detections, axis=0)
model_outputs.update({
'detections': detections,
})
else:
encoded_box_targets = mask_rcnn_architecture.encode_box_targets(
rpn_box_rois, box_targets, class_targets,
params['bbox_reg_weights'])
model_outputs.update({
'rpn_score_outputs': rpn_score_outputs,
'rpn_box_outputs': rpn_box_outputs,
'class_outputs': class_outputs,
'box_outputs': box_outputs,
'class_targets': class_targets,
'box_targets': encoded_box_targets,
'box_rois': rpn_box_rois,
})
# Faster-RCNN mode.
if not params['include_mask']:
return model_outputs
# Mask sampling
if mode != tf.estimator.ModeKeys.TRAIN:
selected_box_rois = detections[:, :, 1:5]
class_indices = tf.to_int32(detections[:, :, 6])
else:
(selected_class_targets, selected_box_targets, selected_box_rois,
proposal_to_label_map) = (
mask_rcnn_architecture.select_fg_for_masks(
class_targets,
box_targets,
rpn_box_rois,
proposal_to_label_map,
max_num_fg=int(params['batch_size_per_im'] *
params['fg_fraction'])))
class_indices = tf.to_int32(selected_class_targets)
mask_roi_features = ops.multilevel_crop_and_resize(fpn_feats,
selected_box_rois,
output_size=14)
mask_outputs = heads.mask_head(
mask_roi_features,
class_indices,
num_classes=params['num_classes'],
mrcnn_resolution=params['mrcnn_resolution'])
model_outputs.update({
'mask_outputs': mask_outputs,
})
if mode == tf.estimator.ModeKeys.TRAIN:
mask_targets = mask_rcnn_architecture.get_mask_targets(
selected_box_rois, proposal_to_label_map, selected_box_targets,
labels['cropped_gt_masks'], params['mrcnn_resolution'])
model_outputs.update({
'mask_targets':
mask_targets,
'selected_class_targets':
selected_class_targets,
})
return model_outputs
batch_size = 64
train_gen = train_datagen.flow(x_train,
to_categorical(y_train),
batch_size=batch_size)
val_gen = val_datagen.flow(x_test,
to_categorical(y_test),
shuffle=False,
batch_size=batch_size)
# Initiate model
tf.keras.backend.clear_session()
if args.model_type == 'resnet':
model, model_with_softmax = resnet_v1([32, 32, 3], 20)
init_weight_copy = model.get_weights()
model_with_softmax.compile(optimizer='adam',
loss='categorical_crossentropy',
metrics=['accuracy'])
elif args.model_type == 'small':
model, model_with_softmax = cifar_small()
init_weight_copy = model.get_weights()
model_with_softmax.compile(optimizer='adam',
loss='categorical_crossentropy',
metrics=['accuracy'])
else:
raise f"No model name {args.model_type}"
def main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
# Enable training summary.
if FLAGS.train_summary_steps > 0:
tf.config.set_soft_device_placement(True)
builder = tfds.builder(FLAGS.dataset, data_dir=FLAGS.data_dir)
builder.download_and_prepare()
num_train_examples = builder.info.splits[FLAGS.train_split].num_examples
num_eval_examples = builder.info.splits[FLAGS.eval_split].num_examples
num_classes = builder.info.features['label'].num_classes
train_steps = model_util.get_train_steps(num_train_examples)
eval_steps = int(math.ceil(num_eval_examples / FLAGS.eval_batch_size))
epoch_steps = int(round(num_train_examples / FLAGS.train_batch_size))
resnet.BATCH_NORM_DECAY = FLAGS.batch_norm_decay
model = resnet.resnet_v1(resnet_depth=FLAGS.resnet_depth,
width_multiplier=FLAGS.width_multiplier,
cifar_stem=FLAGS.image_size <= 32)
checkpoint_steps = (FLAGS.checkpoint_steps
or (FLAGS.checkpoint_epochs * epoch_steps))
cluster = None
if FLAGS.use_tpu and FLAGS.master is None:
if FLAGS.tpu_name:
cluster = tf.distribute.cluster_resolver.TPUClusterResolver(
FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
else:
cluster = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(cluster)
tf.tpu.experimental.initialize_tpu_system(cluster)
default_eval_mode = tf.estimator.tpu.InputPipelineConfig.PER_HOST_V1
sliced_eval_mode = tf.estimator.tpu.InputPipelineConfig.SLICED
run_config = tf.estimator.tpu.RunConfig(
tpu_config=tf.estimator.tpu.TPUConfig(
iterations_per_loop=checkpoint_steps,
eval_training_input_configuration=sliced_eval_mode
if FLAGS.use_tpu else default_eval_mode),
model_dir=FLAGS.model_dir,
save_summary_steps=checkpoint_steps,
save_checkpoints_steps=checkpoint_steps,
keep_checkpoint_max=FLAGS.keep_checkpoint_max,
master=FLAGS.master,
cluster=cluster)
estimator = tf.estimator.tpu.TPUEstimator(
model_lib.build_model_fn(model, num_classes, num_train_examples),
config=run_config,
train_batch_size=FLAGS.train_batch_size,
eval_batch_size=FLAGS.eval_batch_size,
use_tpu=FLAGS.use_tpu)
if FLAGS.mode == 'eval':
for ckpt in tf.train.checkpoints_iterator(run_config.model_dir,
min_interval_secs=15):
try:
result = perform_evaluation(estimator=estimator,
input_fn=data_lib.build_input_fn(
builder, False),
eval_steps=eval_steps,
model=model,
num_classes=num_classes,
checkpoint_path=ckpt)
except tf.errors.NotFoundError:
continue
if result['global_step'] >= train_steps:
return
else:
estimator.train(data_lib.build_input_fn(builder, True),
max_steps=train_steps)
if FLAGS.mode == 'train_then_eval':
perform_evaluation(estimator=estimator,
input_fn=data_lib.build_input_fn(
builder, False),
eval_steps=eval_steps,
model=model,
num_classes=num_classes)
result1 = result.join(label)
def resample(num_samples, level):
if level == 0 or level == 1 or level == 2:
l = result1[result1['level'] == level].index
rnd_idx = random.sample(l, num_samples)
return result1.iloc[rnd_idx, :]
else:
n = len(result1[result1['level'] == level])
l = result1[result1['level'] == level].index
idx = np.random.choice(l, num_samples, n)
return result1.iloc[idx, :]
ls = []
for i in range(5):
ls.append(resample(2000, i))
resampled_df = pd.concat(ls)
resampled_x_train = np.concatenate(
[arr[np.newaxis] for arr in resampled_df.img_arr])
resampled_y_train = resampled_df.iloc[:, 3:].values
from resnet import resnet_v1
model = resnet_v1()
model.fit(resampled_x_train,
resampled_y_train,
validation_split=0.1,
epochs=10,
batch_size=10)
images_per_thread, max_threads)
sk_ratio = 0.0625
input_shape = (224, 224, 3)
width_multiplier = 1
resnet_depth = 50
global_bn = True
batch_norm_decay = 0.9
train_mode = "pretrain"
se_ratio = 0
hidden_layer_sizes = [256, 128, 50]
tf.compat.v1.reset_default_graph()
# base_model = ResNet50(input_shape = input_shape, weights = None)
base_model = resnet_model.resnet_v1(resnet_depth, width_multiplier)
base_model = base_model(Input(shape=input_shape), is_training=True)
base_model.trainable = True
inputs = Input(shape=image_shape)
h = base_model(inputs, training=True)
h = Reshape((1, 1, h.shape[-1]))(h)
h = GlobalAveragePooling2D()(h)
projection_1 = Dense(hidden_layer_sizes[0])(h)
projection_1 = Activation("relu")(projection_1)
projection_2 = Dense(hidden_layer_sizes[1])(projection_1)
projection_2 = Activation("relu")(projection_2)
projection_3 = Dense(hidden_layer_sizes[2])(projection_2)
def main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
# need to import here because we need the above flags
from mobilenetv3.mobilenet_v3 import V3_LARGE_MINIMALISTIC, V3_SMALL_MINIMALISTIC, mobilenet_func
# Enable training summary.
if FLAGS.train_summary_steps > 0:
tf.config.set_soft_device_placement(True)
builder = tfds.builder(FLAGS.dataset, data_dir=FLAGS.data_dir)
builder.download_and_prepare()
num_train_examples = builder.info.splits[FLAGS.train_split].num_examples
num_eval_examples = builder.info.splits[FLAGS.eval_split].num_examples
num_classes = builder.info.features['label'].num_classes
train_steps = model_util.get_train_steps(num_train_examples)
eval_steps = int(math.ceil(num_eval_examples / FLAGS.eval_batch_size))
epoch_steps = int(round(num_train_examples / FLAGS.train_batch_size))
if FLAGS.backbone == "resnet":
resnet.BATCH_NORM_DECAY = FLAGS.batch_norm_decay
model = resnet.resnet_v1(resnet_depth=FLAGS.resnet_depth,
width_multiplier=FLAGS.width_multiplier,
cifar_stem=FLAGS.image_size <= 32,
conv1_stride1=(FLAGS.image_size == 112)
or (FLAGS.image_size == 84))
elif FLAGS.backbone == "mobilenet_v3_large_minimalistic":
model = mobilenet_func(conv_defs=V3_LARGE_MINIMALISTIC,
depth_multiplier=FLAGS.width_multiplier)
elif FLAGS.backbone == "mobilenet_v3_small_minimalistic":
model = mobilenet_func(conv_defs=V3_SMALL_MINIMALISTIC,
depth_multiplier=FLAGS.width_multiplier)
else:
raise ValueError("wrong backbone:" + FLAGS.backbone)
checkpoint_steps = (FLAGS.checkpoint_steps
or (FLAGS.checkpoint_epochs * epoch_steps))
cluster = None
if FLAGS.use_tpu and FLAGS.master is None:
if FLAGS.tpu_name:
cluster = tf.distribute.cluster_resolver.TPUClusterResolver(
FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
else:
cluster = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(cluster)
tf.tpu.experimental.initialize_tpu_system(cluster)
config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))
if FLAGS.use_fp16:
# This works, but no loss rescaling
config.graph_options.rewrite_options.auto_mixed_precision = 1
if FLAGS.use_tpu:
sliced_eval_mode = tf.estimator.tpu.InputPipelineConfig.SLICED
else:
sliced_eval_mode = tf.estimator.tpu.InputPipelineConfig.PER_HOST_V1
run_config = tf.estimator.tpu.RunConfig(
tpu_config=tf.estimator.tpu.TPUConfig(
iterations_per_loop=checkpoint_steps,
eval_training_input_configuration=sliced_eval_mode),
model_dir=FLAGS.model_dir,
save_summary_steps=checkpoint_steps,
save_checkpoints_steps=checkpoint_steps,
keep_checkpoint_max=FLAGS.keep_checkpoint_max,
master=FLAGS.master,
cluster=cluster,
session_config=config)
estimator = tf.estimator.tpu.TPUEstimator(
model_lib.build_model_fn(model, num_classes, num_train_examples),
config=run_config,
train_batch_size=FLAGS.train_batch_size,
eval_batch_size=FLAGS.eval_batch_size,
use_tpu=FLAGS.use_tpu,
model_dir=FLAGS.model_dir)
if FLAGS.mode == 'eval':
for ckpt in tf.train.checkpoints_iterator(run_config.model_dir,
min_interval_secs=15):
try:
result = perform_evaluation(estimator=estimator,
input_fn=data_lib.build_input_fn(
builder, False),
eval_steps=eval_steps,
model=model,
num_classes=num_classes,
checkpoint_path=ckpt)
except tf.errors.NotFoundError:
continue
if result['global_step'] >= train_steps:
return
else:
profile_hook = tf.estimator.ProfilerHook(save_steps=100000000,
save_secs=None,
output_dir=FLAGS.model_dir,
show_dataflow=True,
show_memory=True)
estimator.train(data_lib.build_input_fn(builder, True),
max_steps=train_steps,
hooks=[profile_hook])
if FLAGS.mode == 'train_then_eval':
perform_evaluation(estimator=estimator,
input_fn=data_lib.build_input_fn(
builder, False),
eval_steps=eval_steps,
model=model,
num_classes=num_classes)
def build_model_graph(features, labels, is_training, params):
"""Builds the forward model graph."""
model_outputs = {}
if params['transpose_input'] and is_training:
features['images'] = tf.transpose(features['images'], [3, 0, 1, 2])
batch_size, image_height, image_width, _ = (
features['images'].get_shape().as_list())
if 'source_ids' not in features:
features['source_ids'] = -1 * tf.ones([batch_size], dtype=tf.float32)
all_anchors = anchors.Anchors(params['min_level'], params['max_level'],
params['num_scales'],
params['aspect_ratios'],
params['anchor_scale'],
(image_height, image_width))
with tf.variable_scope('resnet%s' % params['resnet_depth']):
resnet_fn = resnet.resnet_v1(
params['resnet_depth'],
num_batch_norm_group=params['num_batch_norm_group'])
backbone_feats = resnet_fn(features['images'],
(params['is_training_bn'] and is_training))
fpn_feats = fpn.fpn(backbone_feats, params['min_level'],
params['max_level'])
rpn_score_outputs, rpn_box_outputs = heads.rpn_head(
fpn_feats, params['min_level'], params['max_level'],
len(params['aspect_ratios'] * params['num_scales']))
if is_training:
rpn_pre_nms_topn = params['rpn_pre_nms_topn']
rpn_post_nms_topn = params['rpn_post_nms_topn']
else:
rpn_pre_nms_topn = params['test_rpn_pre_nms_topn']
rpn_post_nms_topn = params['test_rpn_post_nms_topn']
rpn_box_scores, rpn_box_rois = roi_ops.multilevel_propose_rois(
rpn_score_outputs,
rpn_box_outputs,
all_anchors,
features['image_info'],
rpn_pre_nms_topn,
rpn_post_nms_topn,
params['rpn_nms_threshold'],
params['rpn_min_size'],
bbox_reg_weights=None,
use_tpu=params['use_tpu'])
rpn_box_rois = tf.to_float(rpn_box_rois)
if is_training:
rpn_box_rois = tf.stop_gradient(rpn_box_rois)
rpn_box_scores = tf.stop_gradient(rpn_box_scores)
if is_training:
# Sampling
box_targets, class_targets, rpn_box_rois, proposal_to_label_map = (
training_ops.proposal_label_op(
rpn_box_rois,
labels['gt_boxes'],
labels['gt_classes'],
features['image_info'],
batch_size_per_im=params['batch_size_per_im'],
fg_fraction=params['fg_fraction'],
fg_thresh=params['fg_thresh'],
bg_thresh_hi=params['bg_thresh_hi'],
bg_thresh_lo=params['bg_thresh_lo']))
# Performs multi-level RoIAlign.
box_roi_features = spatial_transform_ops.multilevel_crop_and_resize(
fpn_feats, rpn_box_rois, output_size=7)
class_outputs, box_outputs, _ = heads.box_head(
box_roi_features,
num_classes=params['num_classes'],
mlp_head_dim=params['fast_rcnn_mlp_head_dim'])
if not is_training:
if params['use_tpu']:
detections = postprocess_ops.generate_detections_tpu(
class_outputs, box_outputs, rpn_box_rois,
features['source_ids'], features['image_info'],
params['test_rpn_post_nms_topn'],
params['test_detections_per_image'], params['test_nms'],
params['bbox_reg_weights'])
else:
detections = postprocess_ops.generate_detections_gpu(
class_outputs, box_outputs, rpn_box_rois,
features['source_ids'], features['image_info'],
params['test_rpn_post_nms_topn'],
params['test_detections_per_image'], params['test_nms'],
params['bbox_reg_weights'])
model_outputs.update({
'detections':
tf.identity(detections, 'Detections'),
})
if params['output_box_features']:
final_box_rois = detections[:, :, 1:5]
final_roi_features = spatial_transform_ops.multilevel_crop_and_resize(
fpn_feats, final_box_rois, output_size=7)
_, _, final_box_features = heads.box_head(
final_roi_features,
num_classes=params['num_classes'],
mlp_head_dim=params['fast_rcnn_mlp_head_dim'])
model_outputs.update({
'box_features':
tf.identity(final_box_features, 'BoxFeatures'),
})
else:
encoded_box_targets = training_ops.encode_box_targets(
rpn_box_rois, box_targets, class_targets,
params['bbox_reg_weights'])
model_outputs.update({
'rpn_score_outputs': rpn_score_outputs,
'rpn_box_outputs': rpn_box_outputs,
'class_outputs': class_outputs,
'box_outputs': box_outputs,
'class_targets': class_targets,
'box_targets': encoded_box_targets,
'box_rois': rpn_box_rois,
})
# Faster-RCNN mode.
if not params['include_mask']:
return model_outputs
# Mask sampling
if not is_training:
selected_box_rois = detections[:, :, 1:5]
class_indices = tf.to_int32(detections[:, :, 6])
else:
(selected_class_targets, selected_box_targets, selected_box_rois,
proposal_to_label_map) = (training_ops.select_fg_for_masks(
class_targets,
box_targets,
rpn_box_rois,
proposal_to_label_map,
max_num_fg=int(params['batch_size_per_im'] *
params['fg_fraction'])))
class_indices = tf.to_int32(selected_class_targets)
mask_roi_features = spatial_transform_ops.multilevel_crop_and_resize(
fpn_feats, selected_box_rois, output_size=14)
mask_outputs = heads.mask_head(mask_roi_features,
class_indices,
num_classes=params['num_classes'],
mrcnn_resolution=params['mrcnn_resolution'])
model_outputs.update({
'mask_outputs': mask_outputs,
})
if is_training:
mask_targets = training_ops.get_mask_targets(
selected_box_rois, proposal_to_label_map, selected_box_targets,
labels['cropped_gt_masks'], params['mrcnn_resolution'])
model_outputs.update({
'mask_targets': mask_targets,
'selected_class_targets': selected_class_targets,
})
else:
model_outputs['mask_outputs'] = tf.identity(
tf.nn.sigmoid(model_outputs['mask_outputs']), 'Masks')
return model_outputs
def main(argv):
#boostx: todo: due to VSC latest code (4-12-2020) coverting
# "--variable_schema='(?!global_step|(?:.*/|^)LARSOptimizer|head)'"
#to "--variable_schema=\'(?!global_step|(?:.*/|^)LARSOptimizer|head)\'"
#I have to use following code to get a workaround: //@audit workaround
FLAGS.variable_schema = '(?!global_step|(?:.*/|^)LARSOptimizer|head)'
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
# Enable training summary.
if FLAGS.train_summary_steps > 0:
tf.config.set_soft_device_placement(True)
#//@follow-up Estmator Evaluation (3)
builder = tfds.builder(FLAGS.dataset, data_dir=FLAGS.data_dir)
builder.download_and_prepare()
num_train_examples = builder.info.splits[FLAGS.train_split].num_examples
num_eval_examples = builder.info.splits[FLAGS.eval_split].num_examples
num_classes = builder.info.features['label'].num_classes
train_steps = model_util.get_train_steps(num_train_examples)
eval_steps = int(math.ceil(num_eval_examples / FLAGS.eval_batch_size))
epoch_steps = int(round(num_train_examples / FLAGS.train_batch_size))
resnet.BATCH_NORM_DECAY = FLAGS.batch_norm_decay
model = resnet.resnet_v1(resnet_depth=FLAGS.resnet_depth,
width_multiplier=FLAGS.width_multiplier,
cifar_stem=FLAGS.image_size <= 32)
checkpoint_steps = (FLAGS.checkpoint_steps
or (FLAGS.checkpoint_epochs * epoch_steps))
master = FLAGS.master
if FLAGS.use_tpu and master is None:
if FLAGS.tpu_name:
cluster = tf.distribute.cluster_resolver.TPUClusterResolver(
FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
else:
cluster = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(cluster)
tf.tpu.experimental.initialize_tpu_system(cluster)
master = cluster.master()
run_config = tf.estimator.tpu.RunConfig(
tpu_config=tf.estimator.tpu.TPUConfig(
iterations_per_loop=checkpoint_steps),
model_dir=FLAGS.model_dir,
save_summary_steps=checkpoint_steps,
save_checkpoints_steps=checkpoint_steps,
keep_checkpoint_max=FLAGS.keep_checkpoint_max,
master=master)
estimator = tf.estimator.tpu.TPUEstimator( #//@follow-up Estmator Evaluation (4)
model_lib.build_model_fn(
model, num_classes,
num_train_examples), #//@follow-up Estmator Evaluation (5)
config=run_config,
train_batch_size=FLAGS.train_batch_size,
eval_batch_size=FLAGS.eval_batch_size,
use_tpu=FLAGS.use_tpu)
if FLAGS.mode == 'eval':
for ckpt in tf.train.checkpoints_iterator(run_config.model_dir,
min_interval_secs=15):
try:
result = perform_evaluation(estimator=estimator,
input_fn=data_lib.build_input_fn(
builder, False),
eval_steps=eval_steps,
model=model,
num_classes=num_classes,
checkpoint_path=ckpt)
except tf.errors.NotFoundError:
continue
if result['global_step'] >= train_steps:
return
else:
estimator.train(data_lib.build_input_fn(builder, True),
max_steps=train_steps)
if FLAGS.mode == 'train_then_eval': #//@follow-up Estmator Evaluation (10)
perform_evaluation(estimator=estimator,
input_fn=data_lib.build_input_fn(
builder, False),
eval_steps=eval_steps,
model=model,
num_classes=num_classes)
def main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
if FLAGS.create_hub:
app.run(create_module_from_checkpoints)
else:
if not os.path.exists(FLAGS.model_dir):
os.makedirs(FLAGS.model_dir)
print("Created directory: {0}".format(os.path.abspath(FLAGS.model_dir)))
# Enable training summary.
if FLAGS.train_summary_steps > 0:
tf.config.set_soft_device_placement(True)
# Choose dataset.
if FLAGS.dataset == "chest_xray":
# Not really a builder, but it's compatible
# TODO config
#data_path = FLAGS.local_tmp_folder
data_path = FLAGS.data_dir
data_split = FLAGS.train_data_split
print(f"***********************************************************************************")
print("")
print(f"DANGER WARNING ON SPLIT -> XRAY Data split:{data_split} SHOULD BE 0.9")
print("")
print(f"***********************************************************************************")
builder, info = chest_xray.XRayDataSet(data_path, config=None, train=True, return_tf_dataset=False, split=data_split)
build_input_fn = partial(data_lib.build_chest_xray_fn, FLAGS.use_multi_gpus, data_path)
num_train_examples = info.get('num_examples')
num_classes = info.get('num_classes')
num_eval_examples = info.get('num_eval_examples')
print(f"num_train_examples:{num_train_examples}, num_eval_examples:{num_eval_examples}")
else:
#builder = tfds.builder(FLAGS.dataset, data_dir=FLAGS.data_dir)
builder.download_and_prepare()
num_train_examples = builder.info.splits[FLAGS.train_split].num_examples
num_eval_examples = builder.info.splits[FLAGS.eval_split].num_examples
num_classes = builder.info.features['label'].num_classes
build_input_fn = data_lib.build_input_fn
train_steps = model_util.get_train_steps(num_train_examples)
eval_steps = int(math.ceil(num_eval_examples / FLAGS.eval_batch_size))
epoch_steps = int(round(num_train_examples / FLAGS.train_batch_size))
resnet.BATCH_NORM_DECAY = FLAGS.batch_norm_decay
model = resnet.resnet_v1(
resnet_depth=FLAGS.resnet_depth,
width_multiplier=FLAGS.width_multiplier,
cifar_stem=FLAGS.image_size <= 32)
checkpoint_steps = (
FLAGS.checkpoint_steps or (FLAGS.checkpoint_epochs * epoch_steps))
cluster = None
if FLAGS.use_tpu and FLAGS.master is None:
if FLAGS.tpu_name:
cluster = tf.distribute.cluster_resolver.TPUClusterResolver(
FLAGS.tpu_name, zone=FLAGS.tpu_zone, project=FLAGS.gcp_project)
else:
cluster = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(cluster)
tf.tpu.experimental.initialize_tpu_system(cluster)
strategy = tf.distribute.MirroredStrategy() if not FLAGS.use_tpu and FLAGS.use_multi_gpus else None # Multi GPU?
print("use_multi_gpus: {0}".format(FLAGS.use_multi_gpus))
print("use MirroredStrategy: {0}".format(not FLAGS.use_tpu and FLAGS.use_multi_gpus))
default_eval_mode = tf.estimator.tpu.InputPipelineConfig.PER_HOST_V1
sliced_eval_mode = tf.estimator.tpu.InputPipelineConfig.SLICED
run_config = tf.estimator.tpu.RunConfig(
tpu_config=tf.estimator.tpu.TPUConfig(
iterations_per_loop=checkpoint_steps,
eval_training_input_configuration=sliced_eval_mode
if FLAGS.use_tpu else default_eval_mode),
model_dir=FLAGS.model_dir,
train_distribute=strategy,
eval_distribute=strategy,
save_summary_steps=checkpoint_steps,
save_checkpoints_steps=checkpoint_steps,
keep_checkpoint_max=FLAGS.keep_checkpoint_max,
master=FLAGS.master,
cluster=cluster)
estimator = tf.estimator.tpu.TPUEstimator(
model_lib.build_model_fn(model, num_classes, num_train_examples, FLAGS.train_batch_size),
config=run_config,
train_batch_size=FLAGS.train_batch_size,
eval_batch_size=FLAGS.eval_batch_size,
use_tpu=FLAGS.use_tpu)
# save flags for this experiment
pickle.dump(FLAGS.flag_values_dict(), open(os.path.join(FLAGS.model_dir, 'experiment_flags.p'), "wb"))
FLAGS.append_flags_into_file(os.path.join(FLAGS.model_dir, 'experiment_flags.txt'))
# Train/Eval
if FLAGS.mode == 'eval':
for ckpt in tf.train.checkpoints_iterator(
run_config.model_dir, min_interval_secs=15):
try:
result = perform_evaluation(
estimator=estimator,
input_fn=build_input_fn(builder, False),
eval_steps=eval_steps,
model=model,
num_classes=num_classes,
checkpoint_path=ckpt)
except tf.errors.NotFoundError:
continue
if result['global_step'] >= train_steps:
return
else:
estimator.train(
build_input_fn(builder, True), max_steps=train_steps)
if FLAGS.mode == 'train_then_eval':
perform_evaluation(
estimator=estimator,
input_fn=build_input_fn(builder, False),
eval_steps=eval_steps,
model=model,
num_classes=num_classes)
# Save the Hub in all case
# app.run(create_module_from_checkpoints)
create_module_from_checkpoints(argv)
# Compute SSL metric:
if FLAGS.compute_ssl_metric:
compute_ssl_metric()
depth = n * 6 + 2
elif version == 2:
depth = n * 9 + 2
# Model name, depth and version
model_type = 'ResNet%dv%d' % (depth, version)
################################################################
# Whenever learning rate reduces, restart the MaSS optimizer at the latest learned weights.
# (If not reducing learning rate during training, one can have only one stage.)
#
#Stage 1: epoch 1-150. Learning rate: 0.1
################################################################
# Build model
model = resnet_v1(input_shape=input_shape, depth=depth)
mass = MaSS(lr=0.1, alpha=0.05, kappa_t=2)
model.compile(loss='categorical_crossentropy',
optimizer=mass,
metrics=['accuracy'])
model.summary()
# Prepare model saving directory.
save_dir = os.path.join(os.getcwd(), 'saved_models')
model_name = 'cifar10_%s_model_mass.{epoch:03d}.h5' % model_type
if not os.path.isdir(save_dir):
os.makedirs(save_dir)
filepath = os.path.join(save_dir, model_name)
checkpoint = ModelCheckpoint(filepath=filepath,
monitor='val_acc',
verbose=1,
def build_model_graph(features, labels, is_training, params):
"""Builds the forward model graph."""
use_batched_nms = (not params['use_tpu'] and params['use_batched_nms'])
is_gpu_inference = (not is_training and use_batched_nms)
model_outputs = {}
if is_training:
if params['transpose_input']:
features['images'] = tf.transpose(features['images'], [2, 0, 1, 3])
batch_size, image_height, image_width, _ = (
features['images'].get_shape().as_list())
# Handles space-to-depth transform.
conv0_space_to_depth_block_size = 0
if is_training:
conv0_space_to_depth_block_size = params[
'conv0_space_to_depth_block_size']
image_height *= conv0_space_to_depth_block_size
image_width *= conv0_space_to_depth_block_size
if 'source_ids' not in features:
features['source_ids'] = -1 * tf.ones([batch_size], dtype=tf.float32)
all_anchors = anchors.Anchors(params['min_level'], params['max_level'],
params['num_scales'],
params['aspect_ratios'],
params['anchor_scale'],
(image_height, image_width))
if 'resnet' in params['backbone']:
with tf.variable_scope(params['backbone']):
resnet_fn = resnet.resnet_v1(
params['backbone'],
conv0_kernel_size=params['conv0_kernel_size'],
conv0_space_to_depth_block_size=conv0_space_to_depth_block_size,
num_batch_norm_group=params['num_batch_norm_group'])
backbone_feats = resnet_fn(
features['images'], (params['is_training_bn'] and is_training))
elif 'mnasnet' in params['backbone']:
with tf.variable_scope(params['backbone']):
_, endpoints = mnasnet_models.build_mnasnet_base(
features['images'],
params['backbone'],
training=(params['is_training_bn'] and is_training),
override_params={'use_keras': False})
backbone_feats = {
2: endpoints['reduction_2'],
3: endpoints['reduction_3'],
4: endpoints['reduction_4'],
5: endpoints['reduction_5'],
}
else:
raise ValueError('Not a valid backbone option: %s' %
params['backbone'])
fpn_feats = fpn.fpn(backbone_feats, params['min_level'],
params['max_level'])
model_outputs.update({
'fpn_features': fpn_feats,
})
rpn_score_outputs, rpn_box_outputs = heads.rpn_head(
fpn_feats, params['min_level'], params['max_level'],
len(params['aspect_ratios'] * params['num_scales']))
if is_training:
rpn_pre_nms_topn = params['rpn_pre_nms_topn']
rpn_post_nms_topn = params['rpn_post_nms_topn']
else:
rpn_pre_nms_topn = params['test_rpn_pre_nms_topn']
rpn_post_nms_topn = params['test_rpn_post_nms_topn']
rpn_box_scores, rpn_box_rois = roi_ops.multilevel_propose_rois(
rpn_score_outputs,
rpn_box_outputs,
all_anchors,
features['image_info'],
rpn_pre_nms_topn,
rpn_post_nms_topn,
params['rpn_nms_threshold'],
params['rpn_min_size'],
bbox_reg_weights=None,
use_batched_nms=use_batched_nms)
rpn_box_rois = tf.to_float(rpn_box_rois)
if is_training:
rpn_box_rois = tf.stop_gradient(rpn_box_rois)
rpn_box_scores = tf.stop_gradient(rpn_box_scores)
if is_training:
# Sampling
box_targets, class_targets, rpn_box_rois, proposal_to_label_map = (
training_ops.proposal_label_op(
rpn_box_rois,
labels['gt_boxes'],
labels['gt_classes'],
features['image_info'],
batch_size_per_im=params['batch_size_per_im'],
fg_fraction=params['fg_fraction'],
fg_thresh=params['fg_thresh'],
bg_thresh_hi=params['bg_thresh_hi'],
bg_thresh_lo=params['bg_thresh_lo']))
# Performs multi-level RoIAlign.
box_roi_features = spatial_transform_ops.multilevel_crop_and_resize(
fpn_feats,
rpn_box_rois,
output_size=7,
is_gpu_inference=is_gpu_inference)
class_outputs, box_outputs, _ = heads.box_head(
box_roi_features,
num_classes=params['num_classes'],
mlp_head_dim=params['fast_rcnn_mlp_head_dim'])
if not is_training:
if is_gpu_inference:
generate_detections_fn = postprocess_ops.generate_detections_gpu
else:
generate_detections_fn = postprocess_ops.generate_detections_tpu
detections = generate_detections_fn(
class_outputs, box_outputs, rpn_box_rois, features['image_info'],
params['test_rpn_post_nms_topn'],
params['test_detections_per_image'], params['test_nms'],
params['bbox_reg_weights'])
model_outputs.update({
'num_detections': detections[0],
'detection_boxes': detections[1],
'detection_classes': detections[2],
'detection_scores': detections[3],
})
else:
encoded_box_targets = training_ops.encode_box_targets(
rpn_box_rois, box_targets, class_targets,
params['bbox_reg_weights'])
model_outputs.update({
'rpn_score_outputs': rpn_score_outputs,
'rpn_box_outputs': rpn_box_outputs,
'class_outputs': class_outputs,
'box_outputs': box_outputs,
'class_targets': class_targets,
'box_targets': encoded_box_targets,
'box_rois': rpn_box_rois,
})
# Faster-RCNN mode.
if not params['include_mask']:
return model_outputs
# Mask sampling
if not is_training:
selected_box_rois = model_outputs['detection_boxes']
class_indices = model_outputs['detection_classes']
# If using GPU for inference, delay the cast until when Gather ops show up
# since GPU inference supports float point better.
# TODO(laigd): revisit this when newer versions of GPU libraries is
# released.
if not is_gpu_inference:
class_indices = tf.to_int32(class_indices)
else:
(selected_class_targets, selected_box_targets, selected_box_rois,
proposal_to_label_map) = (training_ops.select_fg_for_masks(
class_targets,
box_targets,
rpn_box_rois,
proposal_to_label_map,
max_num_fg=int(params['batch_size_per_im'] *
params['fg_fraction'])))
class_indices = tf.to_int32(selected_class_targets)
mask_roi_features = spatial_transform_ops.multilevel_crop_and_resize(
fpn_feats,
selected_box_rois,
output_size=14,
is_gpu_inference=is_gpu_inference)
mask_outputs = heads.mask_head(mask_roi_features,
class_indices,
num_classes=params['num_classes'],
mrcnn_resolution=params['mrcnn_resolution'],
is_gpu_inference=is_gpu_inference)
if is_training:
mask_targets = training_ops.get_mask_targets(
selected_box_rois, proposal_to_label_map, selected_box_targets,
labels['cropped_gt_masks'], params['mrcnn_resolution'])
model_outputs.update({
'mask_outputs': mask_outputs,
'mask_targets': mask_targets,
'selected_class_targets': selected_class_targets,
})
else:
model_outputs.update({
'detection_masks': tf.nn.sigmoid(mask_outputs),
})
return model_outputs
def _model_outputs():
"""Generates outputs from the model."""
model_outputs = {}
with tf.variable_scope('resnet%s' % params['resnet_depth']):
resnet_fn = resnet.resnet_v1(
params['resnet_depth'],
num_batch_norm_group=params['num_batch_norm_group'])
backbone_feats = resnet_fn(features['images'], params['is_training_bn'])
fpn_feats = fpn.fpn(
backbone_feats, params['min_level'], params['max_level'])
rpn_score_outputs, rpn_box_outputs = heads.rpn_head(
fpn_feats,
params['min_level'], params['max_level'],
len(params['aspect_ratios'] * params['num_scales']))
if mode == tf.estimator.ModeKeys.TRAIN:
rpn_pre_nms_topn = params['rpn_pre_nms_topn']
rpn_post_nms_topn = params['rpn_post_nms_topn']
else:
rpn_pre_nms_topn = params['test_rpn_pre_nms_topn']
rpn_post_nms_topn = params['test_rpn_post_nms_topn']
rpn_box_scores, rpn_box_rois = roi_ops.multilevel_propose_rois(
rpn_score_outputs,
rpn_box_outputs,
all_anchors,
features['image_info'],
rpn_pre_nms_topn,
rpn_post_nms_topn,
params['rpn_nms_threshold'],
params['rpn_min_size'],
bbox_reg_weights=None,
use_tpu=params['use_tpu'])
rpn_box_rois = tf.to_float(rpn_box_rois)
if mode == tf.estimator.ModeKeys.TRAIN:
rpn_box_rois = tf.stop_gradient(rpn_box_rois)
rpn_box_scores = tf.stop_gradient(rpn_box_scores)
if mode == tf.estimator.ModeKeys.TRAIN:
# Sampling
box_targets, class_targets, rpn_box_rois, proposal_to_label_map = (
training_ops.proposal_label_op(
rpn_box_rois,
labels['gt_boxes'],
labels['gt_classes'],
features['image_info'],
batch_size_per_im=params['batch_size_per_im'],
fg_fraction=params['fg_fraction'],
fg_thresh=params['fg_thresh'],
bg_thresh_hi=params['bg_thresh_hi'],
bg_thresh_lo=params['bg_thresh_lo']))
# Performs multi-level RoIAlign.
box_roi_features = spatial_transform_ops.multilevel_crop_and_resize(
fpn_feats, rpn_box_rois, output_size=7)
class_outputs, box_outputs, _ = heads.box_head(
box_roi_features, num_classes=params['num_classes'],
mlp_head_dim=params['fast_rcnn_mlp_head_dim'])
if mode != tf.estimator.ModeKeys.TRAIN:
if params['use_tpu']:
detections = postprocess_ops.generate_detections_tpu(
class_outputs,
box_outputs,
rpn_box_rois,
features['source_ids'],
features['image_info'],
params['test_rpn_post_nms_topn'],
params['test_detections_per_image'],
params['test_nms'],
params['bbox_reg_weights'])
else:
detections = postprocess_ops.generate_detections_gpu(
class_outputs,
box_outputs,
rpn_box_rois,
features['source_ids'],
features['image_info'],
params['test_rpn_post_nms_topn'],
params['test_detections_per_image'],
params['test_nms'],
params['bbox_reg_weights'])
detections = tf.identity(detections, 'Detections')
model_outputs.update({
'detections': detections,
})
if params['output_box_features']:
final_box_rois = detections[:, :, 1:5]
final_roi_features = spatial_transform_ops.multilevel_crop_and_resize(
fpn_feats, final_box_rois, output_size=7)
_, _, final_box_features = heads.box_head(
final_roi_features, num_classes=params['num_classes'],
mlp_head_dim=params['fast_rcnn_mlp_head_dim'])
final_box_features = tf.identity(final_box_features, 'BoxFeatures')
model_outputs.update({
'box_features': final_box_features,
})
else:
encoded_box_targets = training_ops.encode_box_targets(
rpn_box_rois, box_targets, class_targets, params['bbox_reg_weights'])
model_outputs.update({
'rpn_score_outputs': rpn_score_outputs,
'rpn_box_outputs': rpn_box_outputs,
'class_outputs': class_outputs,
'box_outputs': box_outputs,
'class_targets': class_targets,
'box_targets': encoded_box_targets,
'box_rois': rpn_box_rois,
})
# Faster-RCNN mode.
if not params['include_mask']:
return model_outputs
# Mask sampling
if mode != tf.estimator.ModeKeys.TRAIN:
selected_box_rois = detections[:, :, 1:5]
class_indices = tf.to_int32(detections[:, :, 6])
else:
(selected_class_targets, selected_box_targets, selected_box_rois,
proposal_to_label_map) = (
training_ops.select_fg_for_masks(
class_targets, box_targets, rpn_box_rois,
proposal_to_label_map,
max_num_fg=int(
params['batch_size_per_im'] * params['fg_fraction'])))
class_indices = tf.to_int32(selected_class_targets)
mask_roi_features = spatial_transform_ops.multilevel_crop_and_resize(
fpn_feats, selected_box_rois, output_size=14)
mask_outputs = heads.mask_head(
mask_roi_features,
class_indices,
num_classes=params['num_classes'],
mrcnn_resolution=params['mrcnn_resolution'])
mask_outputs = tf.identity(mask_outputs, 'Masks')
model_outputs.update({
'mask_outputs': mask_outputs,
})
if mode == tf.estimator.ModeKeys.TRAIN:
mask_targets = training_ops.get_mask_targets(
selected_box_rois, proposal_to_label_map, selected_box_targets,
labels['cropped_gt_masks'], params['mrcnn_resolution'])
model_outputs.update({
'mask_targets': mask_targets,
'selected_class_targets': selected_class_targets,
})
return model_outputs