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_wsl.create(cfg.MODEL.TYPE,
train=False,
gpu_id=gpu_id)
namescope = 'gpu_{}/'.format(gpu_id)
if cfg.WSL.DEEP_MEM:
from detectron.utils.wsl_memonger import share_freeze_blobs
model.net._net = share_freeze_blobs(
model.net,
namescope,
)
net_utils.initialize_gpu_from_weights_file(
model,
weights_file,
gpu_id=gpu_id,
)
model_builder_wsl.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 create_model():
"""Build the model and look for saved model checkpoints in case we can
resume from one.
"""
logger = logging.getLogger(__name__)
start_iter = 0
checkpoints = {}
output_dir = get_output_dir(cfg.TRAIN.DATASETS, training=True)
weights_file = cfg.TRAIN.WEIGHTS
if cfg.TRAIN.AUTO_RESUME:
# Check for the final model (indicates training already finished)
final_path = os.path.join(output_dir, 'model_final.pkl')
if os.path.exists(final_path):
logger.info('model_final.pkl exists; no need to train!')
# return None, None, None, {'final': final_path}, output_dir
files = os.listdir(output_dir)
for f in files:
iter_string = re.findall(r'(?<=model_iter)\d+(?=\.pkl)', f)
if len(iter_string) > 0:
checkpoint_iter = int(iter_string[0])
checkpoints[checkpoint_iter] = os.path.join(output_dir, f)
checkpoints['final'] = final_path
return None, None, None, checkpoints, output_dir
if cfg.TRAIN.COPY_WEIGHTS:
copyfile(
weights_file,
os.path.join(output_dir, os.path.basename(weights_file)))
logger.info('Copy {} to {}'.format(weights_file, output_dir))
# Find the most recent checkpoint (highest iteration number)
files = os.listdir(output_dir)
for f in files:
iter_string = re.findall(r'(?<=model_iter)\d+(?=\.pkl)', f)
if len(iter_string) > 0:
checkpoint_iter = int(iter_string[0])
if checkpoint_iter > start_iter:
# Start one iteration immediately after the checkpoint iter
start_iter = checkpoint_iter + 1
resume_weights_file = f
if start_iter > 0:
# Override the initialization weights with the found checkpoint
weights_file = os.path.join(output_dir, resume_weights_file)
logger.info(
'========> Resuming from checkpoint {} at start iter {}'.
format(weights_file, start_iter)
)
logger.info('Building model: {}'.format(cfg.MODEL.TYPE))
model = model_builder_wsl.create(cfg.MODEL.TYPE, train=True)
if cfg.MEMONGER:
optimize_memory(model)
# Performs random weight initialization as defined by the model
workspace.RunNetOnce(model.param_init_net)
return model, weights_file, start_iter, checkpoints, output_dir
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_wsl.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_wsl.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 create_cpg_net(train=True):
logger = logging.getLogger(__name__)
FREEZE_CONV_BODY = cfg.TRAIN.FREEZE_CONV_BODY
FREEZE_AT = cfg.TRAIN.FREEZE_AT
WSL_CSC = cfg.WSL.CSC
CENTER_LOSS = cfg.WSL.CENTER_LOSS
MIN_ENTROPY_LOSS = cfg.WSL.MIN_ENTROPY_LOSS
MASK_ON = cfg.MODEL.MASK_ON
EXECUTION_TYPE = cfg.MODEL.EXECUTION_TYPE
cfg.immutable(False)
cfg.TRAIN.FREEZE_CONV_BODY = False
cfg.TRAIN.FREEZE_AT = 0
cfg.WSL.CSC = False
cfg.WSL.CPG = False
cfg.WSL.CENTER_LOSS = False
cfg.WSL.MIN_ENTROPY_LOSS = False
cfg.MODEL.MASK_ON = False
cfg.MODEL.EXECUTION_TYPE = b'simple'
cfg.immutable(True)
output_dir = get_output_dir(cfg.TRAIN.DATASETS, training=True)
for gpu_id in range(cfg.NUM_GPUS):
logger.info('Building model: {}'.format('gpu_' + str(gpu_id) + '_' +
cfg.MODEL.TYPE + '_cpg'))
model = model_builder_wsl.create('gpu_' + str(gpu_id) + '_' +
cfg.MODEL.TYPE + '_cpg',
train=train)
# workspace.RunNetOnce(model.param_init_net)
#-----------------------------------------------------------------------------------
# logger.info(
# 'Outputs saved to: {:s}'.format(os.path.abspath(output_dir)))
# dump_proto_files(model, output_dir)
if cfg.MEMONGER and False:
start_op, end_op = OP_surgery_head(model, gpu_id)
optimize_memory_cpg(model, gpu_id)
OP_surgery_back(model, gpu_id, start_op, end_op)
namescope = 'gpu_{}/'.format(gpu_id)
model.net._net.op[0].input[
0] = namescope + cfg.WSL.CPG_PRE_BLOB + '_grad'
model.net._net.op[-1].output[
-1] = namescope + cfg.WSL.CPG_DATA_BLOB + '_grad'
# share_surgery(model, gpu_id)
else:
OP_surgery(model, gpu_id)
Input_surgery(model, gpu_id)
workspace.CreateBlob('gpu_' + str(gpu_id) + '/' +
cfg.WSL.CPG_PRE_BLOB + '_grad')
optimize_memory_cpg(model, gpu_id)
#-----------------------------------------------------------------------------------
# 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)
cfg.immutable(False)
cfg.TRAIN.FREEZE_CONV_BODY = FREEZE_CONV_BODY
cfg.TRAIN.FREEZE_AT = FREEZE_AT
cfg.WSL.CSC = WSL_CSC
cfg.WSL.CENTER_LOSS = CENTER_LOSS
cfg.WSL.MIN_ENTROPY_LOSS = MIN_ENTROPY_LOSS
cfg.MODEL.MASK_ON = MASK_ON
cfg.MODEL.EXECUTION_TYPE = EXECUTION_TYPE
cfg.immutable(True)