if __name__ == '__main__':
#bn = 'v2_openpose+light_maxlikelihood_20191211_150642_adam_lr0.0001_wd0.0_batch32'
#bn = 'v2_openpose+light+fullsym_maxlikelihood_20191229_223132_adam_lr0.0001_wd0.0_batch22'
#bn = 'v2_openpose+head_maxlikelihood_20191223_105436_adam_lr0.0001_wd0.0_batch28'
#bn = 'v2_openpose+light+fullsym_maxlikelihood_20191229_223132_adam_lr0.0001_wd0.0_batch22'
#bn = 'v2_openpose+light+head_maxlikelihood_20191219_144819_adam_lr0.0001_wd0.0_batch32'
#bn = 'v3_openpose+light+head_maxlikelihood_20200118_100732_adam_lr0.0001_wd0.0_batch24'
bn = 'v4PAFflat_openpose+light+head_maxlikelihood_20200206_105708_adam_lr0.0001_wd0.0_batch24'
set_type = bn.partition('_')[0]
model_path = Path.home() / 'workspace/WormData/worm-poses/results' / set_type / bn / 'model_best.pth.tar'
cuda_id = 0
device = get_device(cuda_id)
batch_size = 64
roi_size = 256
images_queue_size = 2
results_queue_size = 4
model_args = get_model_arguments(bn)
#%%
model = PoseDetector(**model_args)
state = torch.load(model_path, map_location = 'cpu')
model.load_state_dict(state['state_dict'])
model.eval()
model = model.to(device)
#%%
def train_PAF(
data_type='v2',
model_name='openpose',
cuda_id=0,
log_dir_root=log_dir_root_dflt,
batch_size=16,
num_workers=1,
loss_type='maxlikelihood',
lr=1e-4,
weight_decay=0.0,
n_epochs=1, #1000,
save_frequency=200,
init_model_path=None):
d_args = data_types[data_type]
flow_args = d_args['flow_args']
root_dir = d_args['root_dir']
log_dir = log_dir_root / data_type
return_bboxes = False
return_half_bboxes = False
if not 'openpose' in model_name:
if 'halfboxes' in data_type:
return_half_bboxes = True
else:
return_bboxes = True
model_args = available_models[model_name]
train_flow = SkelMapsFlow(
root_dir=root_dir,
set2read='train',
#set2read = 'validation',
#samples_per_epoch = 1000,
return_key_value_pairs=True,
PAF_seg_dist=model_args['PAF_seg_dist'],
n_segments=model_args['n_segments'],
fold_skeleton=model_args['fold_skeleton'],
return_bboxes=return_bboxes,
return_half_bboxes=return_half_bboxes,
**flow_args)
val_flow = SkelMapsFlowValidation(
root_dir=root_dir,
set2read='validation',
return_key_value_pairs=True,
PAF_seg_dist=model_args['PAF_seg_dist'],
n_segments=model_args['n_segments'],
fold_skeleton=model_args['fold_skeleton'],
return_bboxes=return_bboxes,
return_half_bboxes=return_half_bboxes,
**flow_args)
if 'openpose' in model_name:
model = PoseDetector(n_segments=train_flow.n_segments_out,
n_affinity_maps=train_flow.n_affinity_maps_out,
n_stages=model_args['n_stages'],
features_type=model_args['features_type'],
use_head_loss=model_args['use_head_loss'],
pose_loss_type=model_args['pose_loss_type'])
else:
model = get_keypointrcnn(backbone=model_args['backbone'],
num_classes=2,
num_keypoints=train_flow.n_segments_out)
if init_model_path is not None:
model_name = 'R+' + model_name
state = torch.load(init_model_path, map_location='cpu')
model.load_state_dict(state['state_dict'])
device = get_device(cuda_id)
lr_scheduler = None
model_params = filter(lambda p: p.requires_grad, model.parameters())
optimizer = torch.optim.Adam(model_params,
lr=lr,
weight_decay=weight_decay)
now = datetime.datetime.now()
date_str = now.strftime('%Y%m%d_%H%M%S')
basename = f'{data_type}_{model_name}_{loss_type}_{date_str}_adam_lr{lr}_wd{weight_decay}_batch{batch_size}'
train_poses(basename,
model,
device,
train_flow,
val_flow,
optimizer,
log_dir,
lr_scheduler=lr_scheduler,
batch_size=batch_size,
n_epochs=n_epochs,
num_workers=num_workers,
save_frequency=save_frequency)