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)