def initialize_model_from_cfg(gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
creates the networks in the Caffe2 workspace.
"""
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
net_utils.initialize_gpu_from_weights_file(
model,
cfg.TEST.WEIGHTS,
gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
if cfg.MODEL.MASK_ON:
workspace.CreateNet(model.mask_net)
if cfg.MODEL.KEYPOINTS_ON:
workspace.CreateNet(model.keypoint_net)
graph = net_drawer.GetPydotGraphMinimal(model.net.Proto().op,
rankdir="LR",
minimal_dependency=True)
#display.Image(graph.create_png(), height=1800)
graph.write_png('graph.png')
with open(os.path.join("video/keypoint", "net.pbtxt"), 'w') as fid:
fid.write(str(model.net.Proto()))
return model
def test_restore_checkpoint():
# Create Model
model = model_builder.create(cfg.MODEL.TYPE, train=True)
add_momentum_init_ops(model)
init_weights(model)
# Fill input blobs
roidb = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES
)
model_builder.add_training_inputs(model, roidb=roidb)
workspace.CreateNet(model.net)
# Bookkeeping for checkpoint creation
iter_num = 0
checkpoints = {}
output_dir = get_output_dir(cfg.TRAIN.DATASETS, training=True)
chk_file_path = os.path.join(output_dir, 'model_iter{}.pkl'.format(iter_num))
checkpoints[iter_num] = chk_file_path
# Save model weights
nu.save_model_to_weights_file(checkpoints[iter_num], model)
orig_gpu_0_params, orig_all_params = get_params(model)
# Change the model weights
init_weights(model)
# Reload the weights in the model
nu.initialize_gpu_from_weights_file(model, chk_file_path, gpu_id=0)
nu.broadcast_parameters(model)
shutil.rmtree(cfg.OUTPUT_DIR)
_, restored_all_params = get_params(model)
# Check if all params are loaded correctly
for scoped_name, blob in orig_all_params.items():
np.testing.assert_array_equal(blob, restored_all_params[scoped_name])
# Check if broadcast_parameters works
for scoped_name, blob in restored_all_params.items():
unscoped_name = c2_utils.UnscopeName(scoped_name)
np.testing.assert_array_equal(blob, orig_gpu_0_params[unscoped_name])
def initialize_model_from_cfg(weights_file, gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
creates the networks in the Caffe2 workspace.
"""
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
net_utils.initialize_gpu_from_weights_file(
model, weights_file, gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
return model
def initialize_model_from_cfg(gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
creates the networks in the Caffe2 workspace.
"""
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
net_utils.initialize_gpu_from_weights_file(
model, cfg.TEST.WEIGHTS, gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
if cfg.MODEL.MASK_ON:
workspace.CreateNet(model.mask_net)
if cfg.MODEL.KEYPOINTS_ON:
workspace.CreateNet(model.keypoint_net)
return model
def initialize_model_from_cfg(gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
creates the networks in the Caffe2 workspace.
"""
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
net_utils.initialize_gpu_from_weights_file(
model, cfg.TEST.WEIGHTS, gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
if cfg.MODEL.MASK_ON:
workspace.CreateNet(model.mask_net)
if cfg.MODEL.KEYPOINTS_ON:
workspace.CreateNet(model.keypoint_net)
return model
def generate_rpn_on_range(
weights_file,
dataset_name,
_proposal_file_ignored,
output_dir,
ind_range=None,
gpu_id=0
):
"""Run inference on all images in a dataset or over an index range of images
in a dataset using a single GPU.
"""
assert cfg.MODEL.RPN_ONLY or cfg.MODEL.FASTER_RCNN
roidb, start_ind, end_ind, total_num_images = get_roidb(
dataset_name, ind_range
)
logger.info(
'Output will be saved to: {:s}'.format(os.path.abspath(output_dir))
)
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
nu.initialize_gpu_from_weights_file(
model, weights_file, gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
boxes, scores, ids = generate_proposals_on_roidb(
model,
roidb,
start_ind=start_ind,
end_ind=end_ind,
total_num_images=total_num_images,
gpu_id=gpu_id,
)
cfg_yaml = yaml.dump(cfg)
if ind_range is not None:
rpn_name = 'rpn_proposals_range_%s_%s.pkl' % tuple(ind_range)
else:
rpn_name = 'rpn_proposals.pkl'
rpn_file = os.path.join(output_dir, rpn_name)
save_object(
dict(boxes=boxes, scores=scores, ids=ids, cfg=cfg_yaml), rpn_file
)
logger.info('Wrote RPN proposals to {}'.format(os.path.abspath(rpn_file)))
return boxes, scores, ids, rpn_file
def generate_rpn_on_range(weights_file,
dataset_name,
_proposal_file_ignored,
output_dir,
ind_range=None,
gpu_id=0):
"""Run inference on all images in a dataset or over an index range of images
in a dataset using a single GPU.
"""
assert cfg.MODEL.RPN_ONLY or cfg.MODEL.FASTER_RCNN
roidb, start_ind, end_ind, total_num_images = get_roidb(
dataset_name, ind_range)
logger.info('Output will be saved to: {:s}'.format(
os.path.abspath(output_dir)))
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
nu.initialize_gpu_from_weights_file(
model,
weights_file,
gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
boxes, scores, ids = generate_proposals_on_roidb(
model,
roidb,
start_ind=start_ind,
end_ind=end_ind,
total_num_images=total_num_images,
gpu_id=gpu_id,
)
cfg_yaml = yaml.dump(cfg)
if ind_range is not None:
rpn_name = 'rpn_proposals_range_%s_%s.pkl' % tuple(ind_range)
else:
rpn_name = 'rpn_proposals.pkl'
rpn_file = os.path.join(output_dir, rpn_name)
save_object(dict(boxes=boxes, scores=scores, ids=ids, cfg=cfg_yaml),
rpn_file)
logger.info('Wrote RPN proposals to {}'.format(os.path.abspath(rpn_file)))
return boxes, scores, ids, rpn_file
def setup_model_for_training(model, output_dir):
"""Loaded saved weights and create the network in the C2 workspace."""
logger = logging.getLogger(__name__)
add_model_training_inputs(model)
if cfg.TRAIN.WEIGHTS:
# Override random weight initialization with weights from a saved model
nu.initialize_gpu_from_weights_file(model, cfg.TRAIN.WEIGHTS, gpu_id=0)
# Even if we're randomly initializing we still need to synchronize
# parameters across GPUs
nu.broadcast_parameters(model)
workspace.CreateNet(model.net)
logger.info('Outputs saved to: {:s}'.format(os.path.abspath(output_dir)))
dump_proto_files(model, output_dir)
# Start loading mini-batches and enqueuing blobs
model.roi_data_loader.register_sigint_handler()
model.roi_data_loader.start(prefill=True)
return output_dir
def setup_model_for_training(model, output_dir):
"""Loaded saved weights and create the network in the C2 workspace."""
logger = logging.getLogger(__name__)
add_model_training_inputs(model)
if cfg.TRAIN.WEIGHTS:
# Override random weight initialization with weights from a saved model
nu.initialize_gpu_from_weights_file(model, cfg.TRAIN.WEIGHTS, gpu_id=0)
# Even if we're randomly initializing we still need to synchronize
# parameters across GPUs
nu.broadcast_parameters(model)
workspace.CreateNet(model.net)
logger.info('Outputs saved to: {:s}'.format(os.path.abspath(output_dir)))
dump_proto_files(model, output_dir)
# Start loading mini-batches and enqueuing blobs
model.roi_data_loader.register_sigint_handler()
model.roi_data_loader.start(prefill=True)
return output_dir
