def val():
if args.is_distributed:
init()
args.rank = get_rank()
args.group_size = get_group_size()
if not os.path.exists(args.output_path):
os.mkdir(args.output_path)
network = OpenPoseNet(vgg_with_bn=params['vgg_with_bn'])
network.set_train(False)
load_model(network, args.model_path)
print("load models right")
dataset = valdata(args.ann,
args.imgpath_val,
args.rank,
args.group_size,
mode='val')
dataset_size = dataset.get_dataset_size()
de_dataset = dataset.create_tuple_iterator()
print("eval dataset size: ", dataset_size)
kpt_json = []
for _, (img, img_id) in tqdm(enumerate(de_dataset), total=dataset_size):
img = img.asnumpy()
img_id = int((img_id.asnumpy())[0])
poses, scores = detect(img, network)
if poses.shape[0] > 0:
#print("got poses")
for index, pose in enumerate(poses):
data = dict()
pose = pose[[
0, 15, 14, 17, 16, 5, 2, 6, 3, 7, 4, 11, 8, 12, 9, 13, 10,
1
], :].round().astype('i')
keypoints = pose.reshape(-1).tolist()
keypoints = keypoints[:-3]
data['image_id'] = img_id
data['score'] = scores[index]
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
else:
print("Predict poses size is zero.", flush=True)
img = draw_person_pose(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), poses)
#print('Saving result into',str(img_id)+'.png...')
save_path = os.path.join(args.output_path, str(img_id) + ".png")
cv2.imwrite(save_path, img)
result_json = 'eval_result.json'
with open(os.path.join(args.output_path, result_json), 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP(os.path.join(args.output_path, result_json),
ann_file=args.ann)
print('result: ', res)
def train():
"""Train function."""
args.outputs_dir = params['save_model_path']
if args.group_size > 1:
init()
context.set_auto_parallel_context(
device_num=get_group_size(),
parallel_mode=ParallelMode.DATA_PARALLEL,
gradients_mean=True)
args.outputs_dir = os.path.join(args.outputs_dir,
"ckpt_{}/".format(str(get_rank())))
args.rank = get_rank()
else:
args.outputs_dir = os.path.join(args.outputs_dir, "ckpt_0/")
args.rank = 0
if args.group_size > 1:
args.max_epoch = params["max_epoch_train_NP"]
args.loss_scale = params['loss_scale'] / 2
args.lr_steps = list(map(int, params["lr_steps_NP"].split(',')))
params['train_type'] = params['train_type_NP']
params['optimizer'] = params['optimizer_NP']
params['group_params'] = params['group_params_NP']
else:
args.max_epoch = params["max_epoch_train"]
args.loss_scale = params['loss_scale']
args.lr_steps = list(map(int, params["lr_steps"].split(',')))
# create network
print('start create network')
criterion = openpose_loss()
criterion.add_flags_recursive(fp32=True)
network = OpenPoseNet(vggpath=params['vgg_path'],
vgg_with_bn=params['vgg_with_bn'])
if params["load_pretrain"]:
print("load pretrain model:", params["pretrained_model_path"])
load_model(network, params["pretrained_model_path"])
train_net = BuildTrainNetwork(network, criterion)
# create dataset
if os.path.exists(args.jsonpath_train) and os.path.exists(args.imgpath_train) \
and os.path.exists(args.maskpath_train):
print('start create dataset')
else:
print('Error: wrong data path')
return 0
num_worker = 20 if args.group_size > 1 else 48
de_dataset_train = create_dataset(args.jsonpath_train,
args.imgpath_train,
args.maskpath_train,
batch_size=params['batch_size'],
rank=args.rank,
group_size=args.group_size,
num_worker=num_worker,
multiprocessing=True,
shuffle=True,
repeat_num=1)
steps_per_epoch = de_dataset_train.get_dataset_size()
print("steps_per_epoch: ", steps_per_epoch)
# lr scheduler
lr_stage, lr_base, lr_vgg = get_lr(params['lr'] * args.group_size,
params['lr_gamma'],
steps_per_epoch,
args.max_epoch,
args.lr_steps,
args.group_size,
lr_type=params['lr_type'],
warmup_epoch=params['warmup_epoch'])
# optimizer
if params['group_params']:
vgg19_base_params = list(
filter(lambda x: 'base.vgg_base' in x.name,
train_net.trainable_params()))
base_params = list(
filter(lambda x: 'base.conv' in x.name,
train_net.trainable_params()))
stages_params = list(
filter(lambda x: 'base' not in x.name,
train_net.trainable_params()))
group_params = [{
'params': vgg19_base_params,
'lr': lr_vgg
}, {
'params': base_params,
'lr': lr_base
}, {
'params': stages_params,
'lr': lr_stage
}]
if params['optimizer'] == "Momentum":
opt = Momentum(group_params, learning_rate=lr_stage, momentum=0.9)
elif params['optimizer'] == "Adam":
opt = Adam(group_params)
else:
raise ValueError("optimizer not support.")
else:
if params['optimizer'] == "Momentum":
opt = Momentum(train_net.trainable_params(),
learning_rate=lr_stage,
momentum=0.9)
elif params['optimizer'] == "Adam":
opt = Adam(train_net.trainable_params(), learning_rate=lr_stage)
else:
raise ValueError("optimizer not support.")
# callback
config_ck = CheckpointConfig(
save_checkpoint_steps=params['ckpt_interval'],
keep_checkpoint_max=params["keep_checkpoint_max"])
ckpoint_cb = ModelCheckpoint(prefix='{}'.format(args.rank),
directory=args.outputs_dir,
config=config_ck)
time_cb = TimeMonitor(data_size=de_dataset_train.get_dataset_size())
if args.rank == 0:
callback_list = [MyLossMonitor(), time_cb, ckpoint_cb]
else:
callback_list = [MyLossMonitor(), time_cb]
# train
if params['train_type'] == 'clip_grad':
train_net = TrainOneStepWithClipGradientCell(train_net,
opt,
sens=args.loss_scale)
train_net.set_train()
model = Model(train_net)
elif params['train_type'] == 'fix_loss_scale':
loss_scale_manager = FixedLossScaleManager(args.loss_scale,
drop_overflow_update=False)
train_net.set_train()
model = Model(train_net,
optimizer=opt,
loss_scale_manager=loss_scale_manager)
else:
raise ValueError("Type {} is not support.".format(
params['train_type']))
print("============== Starting Training ==============")
model.train(args.max_epoch,
de_dataset_train,
callbacks=callback_list,
dataset_sink_mode=False)
return 0
def create_network(name, *args, **kwargs):
if name == "openposenet":
return OpenPoseNet(*args, **kwargs)
raise NotImplementedError(f"{name} is not implemented in the repo")
def openposenet(*args, **kwargs):
return OpenPoseNet(*args, **kwargs)
parser.add_argument("--file_format",
type=str,
choices=["AIR", "ONNX", "MINDIR"],
default="AIR",
help="file format")
parser.add_argument("--device_target",
type=str,
default="Ascend",
choices=["Ascend", "GPU", "CPU"],
help="device target (default: Ascend)")
args = parser.parse_args()
context.set_context(mode=context.GRAPH_MODE,
device_target=args.device_target,
device_id=args.device_id)
if __name__ == "__main__":
context.set_context(mode=context.GRAPH_MODE, save_graphs=False)
# define net
net = OpenPoseNet()
# load checkpoint
param_dict = load_checkpoint(args.ckpt_file)
load_param_into_net(net, param_dict)
inputs = np.ones([args.batch_size, 3, params["insize"],
params["insize"]]).astype(np.float32)
export(net,
Tensor(inputs),
file_name=args.file_name,
file_format=args.file_format)
def train():
"""Train function."""
args = parse_args()
args.outputs_dir = params['save_model_path']
if args.group_size > 1:
init()
context.set_auto_parallel_context(device_num=get_group_size(), parallel_mode=ParallelMode.DATA_PARALLEL,
gradients_mean=True)
args.outputs_dir = os.path.join(args.outputs_dir, "ckpt_{}/".format(str(get_rank())))
args.rank = get_rank()
else:
args.outputs_dir = os.path.join(args.outputs_dir, "ckpt_0/")
args.rank = 0
# with out loss_scale
if args.group_size > 1:
args.loss_scale = params['loss_scale'] / 2
args.lr_steps = list(map(int, params["lr_steps_NP"].split(',')))
else:
args.loss_scale = params['loss_scale']
args.lr_steps = list(map(int, params["lr_steps"].split(',')))
# create network
print('start create network')
criterion = openpose_loss()
criterion.add_flags_recursive(fp32=True)
network = OpenPoseNet(vggpath=params['vgg_path'])
# network.add_flags_recursive(fp32=True)
if params["load_pretrain"]:
print("load pretrain model:", params["pretrained_model_path"])
load_model(network, params["pretrained_model_path"])
train_net = BuildTrainNetwork(network, criterion)
# create dataset
if os.path.exists(args.jsonpath_train) and os.path.exists(args.imgpath_train) \
and os.path.exists(args.maskpath_train):
print('start create dataset')
else:
print('Error: wrong data path')
num_worker = 20 if args.group_size > 1 else 48
de_dataset_train = create_dataset(args.jsonpath_train, args.imgpath_train, args.maskpath_train,
batch_size=params['batch_size'],
rank=args.rank,
group_size=args.group_size,
num_worker=num_worker,
multiprocessing=True,
shuffle=True,
repeat_num=1)
steps_per_epoch = de_dataset_train.get_dataset_size()
print("steps_per_epoch: ", steps_per_epoch)
# lr scheduler
lr_stage, lr_base, lr_vgg = get_lr(params['lr'] * args.group_size,
params['lr_gamma'],
steps_per_epoch,
params["max_epoch_train"],
args.lr_steps,
args.group_size)
vgg19_base_params = list(filter(lambda x: 'base.vgg_base' in x.name, train_net.trainable_params()))
base_params = list(filter(lambda x: 'base.conv' in x.name, train_net.trainable_params()))
stages_params = list(filter(lambda x: 'base' not in x.name, train_net.trainable_params()))
group_params = [{'params': vgg19_base_params, 'lr': lr_vgg},
{'params': base_params, 'lr': lr_base},
{'params': stages_params, 'lr': lr_stage}]
opt = Adam(group_params, loss_scale=args.loss_scale)
train_net.set_train(True)
loss_scale_manager = FixedLossScaleManager(args.loss_scale, drop_overflow_update=False)
model = Model(train_net, optimizer=opt, loss_scale_manager=loss_scale_manager)
params['ckpt_interval'] = max(steps_per_epoch, params['ckpt_interval'])
config_ck = CheckpointConfig(save_checkpoint_steps=params['ckpt_interval'],
keep_checkpoint_max=params["keep_checkpoint_max"])
ckpoint_cb = ModelCheckpoint(prefix='{}'.format(args.rank), directory=args.outputs_dir, config=config_ck)
time_cb = TimeMonitor(data_size=de_dataset_train.get_dataset_size())
callback_list = [MyLossMonitor(), time_cb, ckpoint_cb]
print("============== Starting Training ==============")
model.train(params["max_epoch_train"], de_dataset_train, callbacks=callback_list,
dataset_sink_mode=False)