m53297601
/
mindspore
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

openpose network optimize 8p acc.

Browse Source
pull/10092/head
zhanghuiyao 5 years ago
parent da74482812
commit 3c5a8538d5
11 changed files with 420 additions and 354 deletions
Show Stats Download Patch File Download Diff File

16
model_zoo/official/cv/openpose/README.md
Unescape View File

@ -142,10 +142,10 @@ Parameters for both training and evaluation can be set in config.py
'batch_size': 10 # training batch size
'lr_gamma': 0.1 # lr scale when reach lr_steps
'lr_steps': '100000,200000,250000' # the steps when lr * lr_gamma
'loss scale': 16386 # the loss scale of mixed precision
'loss scale': 16384 # the loss scale of mixed precision
'max_epoch_train': 60 # total training epochs
'insize': 368 # image size used as input to the model
'keep_checkpoint_max': 5 # only keep the last keep_checkpoint_max checkpoint
'keep_checkpoint_max': 1 # only keep the last keep_checkpoint_max checkpoint
'log_interval': 100 # the interval of print a log
'ckpt_interval': 5000 # the interval of saving a output model
```
@ -195,7 +195,7 @@ For more configuration details, please refer the script `config.py`.
```python
# grep "AP" eval.log
{'AP': 0.40030956300341397, 'Ap .5': 0.6658941566481336, 'AP .75': 0.396047897339743, 'AP (M)': 0.3075356543635785, 'AP (L)': 0.533772768618845, 'AR': 0.4519836272040302, 'AR .5': 0.693639798488665, 'AR .75': 0.4570214105793451, 'AR (M)': 0.32155148866429945, 'AR (L)': 0.6330360460795242}
{'AP': 0.39830956300341397, 'Ap .5': 0.6658941566481336, 'AP .75': 0.396047897339743, 'AP (M)': 0.3075356543635785, 'AP (L)': 0.533772768618845, 'AR': 0.4519836272040302, 'AR .5': 0.693639798488665, 'AR .75': 0.4570214105793451, 'AR (M)': 0.32155148866429945, 'AR (L)': 0.6330360460795242}
```
@ -209,14 +209,14 @@ For more configuration details, please refer the script `config.py`.
| -------------------------- | -----------------------------------------------------------
| Model Version | openpose
| Resource | Ascend 910 ；CPU 2.60GHz，192cores；Memory，755G
| uploaded Date | 10/20/2020 (month/day/year)
| uploaded Date | 12/14/2020 (month/day/year)
| MindSpore Version | 1.0.1-alpha
| Training Parameters | epoch = 60, steps = 30k, batch_size = 10, lr = 0.0001
| Optimizer | Adam
| Training Parameters | epoch=60(1pcs)/80(8pcs), steps=30k(1pcs)/5k(8pcs), batch_size=10, init_lr=0.0001
| Optimizer | Adam(1pcs)/Momentum(8pcs)
| Loss Function | MSE
| outputs | pose
| Speed | 1pc: 29imgs/s
| Total time | 1pc: 30h
| Speed | 1pcs: 35fps, 8pcs: 230fps
| Total time | 1pcs: 22.5h, 8pcs: 5.1h
| Checkpoint for Fine tuning | 602.33M (.ckpt file)

83
model_zoo/official/cv/openpose/eval.py
Unescape View File

@ -17,10 +17,9 @@ import os
import argparse
import warnings
import sys
import cv2
from tqdm import tqdm
import numpy as np
from tqdm import tqdm
import cv2
from scipy.ndimage.filters import gaussian_filter
from pycocotools.coco import COCO as LoadAnn
from pycocotools.cocoeval import COCOeval as MapEval
@ -30,9 +29,10 @@ from mindspore.train.serialization import load_checkpoint, load_param_into_net
from mindspore.communication.management import init, get_rank, get_group_size
from mindspore.common import dtype as mstype
from src.dataset import valdata
from src.openposenet import OpenPoseNet
from src.config import params, JointType
from src.openposenet import OpenPoseNet
from src.dataset import valdata
warnings.filterwarnings("ignore")
devid = int(os.getenv('DEVICE_ID'))
@ -40,6 +40,18 @@ context.set_context(mode=context.GRAPH_MODE,
device_target="Ascend", save_graphs=False, device_id=devid)
show_gt = 0
parser = argparse.ArgumentParser('mindspore openpose_net test')
parser.add_argument('--model_path', type=str, default='./0-33_170000.ckpt', help='path of testing model')
parser.add_argument('--imgpath_val', type=str, default='./dataset/coco/val2017', help='path of testing imgs')
parser.add_argument('--ann', type=str, default='./dataset/coco/annotations/person_keypoints_val2017.json',
help='path of annotations')
parser.add_argument('--output_path', type=str, default='./output_img', help='path of testing imgs')
# distributed related
parser.add_argument('--is_distributed', type=int, default=0, help='if multi device')
parser.add_argument('--rank', type=int, default=0, help='local rank of distributed')
parser.add_argument('--group_size', type=int, default=1, help='world size of distributed')
args, _ = parser.parse_known_args()
def evaluate_mAP(res_file, ann_file, ann_type='keypoints', silence=True):
class NullWriter():
def write(self, arg):
@ -68,23 +80,6 @@ def evaluate_mAP(res_file, ann_file, ann_type='keypoints', silence=True):
return info_str
def parse_args():
"""Parse arguments."""
parser = argparse.ArgumentParser('mindspore openpose_net test')
parser.add_argument('--model_path', type=str, default='./scripts/train_parallel0/checkpoints/ckpt_0/0-60_663.ckpt',
help='path of testing model')
parser.add_argument('--imgpath_val', type=str, default='/data0/zhy/dataset/coco/val2017',
help='path of testing imgs')
parser.add_argument('--ann', type=str, default='/data0/zhy/dataset/coco/annotations/person_keypoints_val2017.json',
help='path of annotations')
parser.add_argument('--output_path', type=str, default='./output_img', help='path of testing imgs')
# distributed related
parser.add_argument('--is_distributed', type=int, default=0, help='if multi device')
parser.add_argument('--rank', type=int, default=0, help='local rank of distributed')
parser.add_argument('--group_size', type=int, default=1, help='world size of distributed')
args, _ = parser.parse_known_args()
return args
def load_model(test_net, model_path):
assert os.path.exists(model_path)
@ -178,7 +173,7 @@ def compute_peaks_from_heatmaps(heatmaps):
return all_peaks
def compute_candidate_connections(paf, cand_a, cand_b, img_len, cfg):
def compute_candidate_connections(paf, cand_a, cand_b, img_len, params_):
candidate_connections = []
for joint_a in cand_a:
for joint_b in cand_b:
@ -186,33 +181,33 @@ def compute_candidate_connections(paf, cand_a, cand_b, img_len, cfg):
norm = np.linalg.norm(vector)
if norm == 0:
continue
ys = np.linspace(joint_a[1], joint_b[1], num=cfg['n_integ_points'])
xs = np.linspace(joint_a[0], joint_b[0], num=cfg['n_integ_points'])
ys = np.linspace(joint_a[1], joint_b[1], num=params_['n_integ_points'])
xs = np.linspace(joint_a[0], joint_b[0], num=params_['n_integ_points'])
integ_points = np.stack([ys, xs]).T.round().astype('i')
paf_in_edge = np.hstack([paf[0][np.hsplit(integ_points, 2)], paf[1][np.hsplit(integ_points, 2)]])
unit_vector = vector / norm
inner_products = np.dot(paf_in_edge, unit_vector)
integ_value = inner_products.sum() / len(inner_products)
integ_value_with_dist_prior = integ_value + min(cfg['limb_length_ratio'] * img_len / norm -
cfg['length_penalty_value'], 0)
n_valid_points = sum(inner_products > cfg['inner_product_thresh'])
if n_valid_points > cfg['n_integ_points_thresh'] and integ_value_with_dist_prior > 0:
integ_value_with_dist_prior = integ_value + min(params_['limb_length_ratio'] * img_len / norm -
params_['length_penalty_value'], 0)
n_valid_points = sum(inner_products > params_['inner_product_thresh'])
if n_valid_points > params_['n_integ_points_thresh'] and integ_value_with_dist_prior > 0:
candidate_connections.append([int(joint_a[3]), int(joint_b[3]), integ_value_with_dist_prior])
candidate_connections = sorted(candidate_connections, key=lambda x: x[2], reverse=True)
return candidate_connections
def compute_connections(pafs, all_peaks, img_len, cfg):
def compute_connections(pafs, all_peaks, img_len, params_):
all_connections = []
for i in range(len(cfg['limbs_point'])):
for i in range(len(params_['limbs_point'])):
paf_index = [i * 2, i * 2 + 1]
paf = pafs[paf_index] # shape: (2, 320, 320)
limb_point = cfg['limbs_point'][i] # example: [<JointType.Neck: 1>, <JointType.RightWaist: 8>]
limb_point = params_['limbs_point'][i] # example: [<JointType.Neck: 1>, <JointType.RightWaist: 8>]
cand_a = all_peaks[all_peaks[:, 0] == limb_point[0]][:, 1:]
cand_b = all_peaks[all_peaks[:, 0] == limb_point[1]][:, 1:]
if cand_a.shape[0] > 0 and cand_b.shape[0] > 0:
candidate_connections = compute_candidate_connections(paf, cand_a, cand_b, img_len, cfg)
candidate_connections = compute_candidate_connections(paf, cand_a, cand_b, img_len, params_)
connections = np.zeros((0, 3))
@ -226,11 +221,11 @@ def compute_connections(pafs, all_peaks, img_len, cfg):
all_connections.append(np.zeros((0, 3)))
return all_connections
def grouping_key_points(all_connections, candidate_peaks, cfg):
def grouping_key_points(all_connections, candidate_peaks, params_):
subsets = -1 * np.ones((0, 20))
for l, connections in enumerate(all_connections):
joint_a, joint_b = cfg['limbs_point'][l]
joint_a, joint_b = params_['limbs_point'][l]
for ind_a, ind_b, score in connections[:, :3]:
ind_a, ind_b = int(ind_a), int(ind_b)
joint_found_cnt = 0
@ -249,6 +244,7 @@ def grouping_key_points(all_connections, candidate_peaks, cfg):
found_subset[-1] += 1 # increment joint count
found_subset[-2] += candidate_peaks[ind_b, 3] + score
elif joint_found_cnt == 2:
found_subset_1 = subsets[joint_found_subset_index[0]]
@ -289,10 +285,8 @@ def grouping_key_points(all_connections, candidate_peaks, cfg):
pass
# delete low score subsets
keep = np.logical_and(subsets[:, -1] >= cfg['n_subset_limbs_thresh'],
subsets[:, -2] / subsets[:, -1] >= cfg['subset_score_thresh'])
# cfg['n_subset_limbs_thresh'] = 3
# cfg['subset_score_thresh'] = 0.2
keep = np.logical_and(subsets[:, -1] >= params_['n_subset_limbs_thresh'],
subsets[:, -2] / subsets[:, -1] >= params_['subset_score_thresh'])
subsets = subsets[keep]
return subsets
@ -319,7 +313,7 @@ def detect(img, network):
# map_w, map_h = compute_optimal_size(orig_img, params['heatmap_size']) # 320
map_w, map_h = compute_optimal_size(orig_img, params['inference_img_size'])
print("image size is: ", input_w, input_h)
# print("image size is: ", input_w, input_h)
resized_image = cv2.resize(orig_img, (input_w, input_h))
x_data = preprocess(resized_image)
@ -394,7 +388,7 @@ def draw_person_pose(orig_img, poses):
return canvas
def depreprocess(img):
# x_data = img.astype('f')
#x_data = img.astype('f')
x_data = img[0]
x_data += 0.5
x_data *= 255
@ -402,15 +396,14 @@ def depreprocess(img):
x_data = x_data.transpose(1, 2, 0)
return x_data
def _eval():
args = parse_args()
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()
network = OpenPoseNet(vgg_with_bn=params['vgg_with_bn'])
network.set_train(False)
load_model(network, args.model_path)
@ -455,4 +448,4 @@ def _eval():
print('result: ', res)
if __name__ == "__main__":
_eval()
val()

7
model_zoo/official/cv/openpose/scripts/run_eval_ascend.sh
Unescape View File

@ -15,9 +15,8 @@
# ============================================================================
export DEVICE_ID=0
export RANK_ID=0
python eval.py \
--model_path ./scripts/train_parallel0/checkpoints/ckpt_0/0-60_663.ckpt \
--imgpath_val /data0/zhy/dataset/coco/val2017 \
--ann /data0/zhy/dataset/coco/annotations/person_keypoints_val2017.json \
--model_path ./scripts/train_parallel0/checkpoints/ckpt_0/0-80_663.ckpt \
--imgpath_val ./dataset/val2017 \
--ann ./dataset/annotations/person_keypoints_val2017.json \
> eval.log 2>&1 &

1
model_zoo/official/cv/openpose/scripts/run_standalone_train.sh
Unescape View File

@ -14,5 +14,6 @@
# limitations under the License.
# ============================================================================
export DEVICE_ID=0
cd ..
python train.py --train_dir train2017 --train_ann person_keypoints_train2017.json > scripts/train.log 2>&1 &

42
model_zoo/official/cv/openpose/src/config.py
Unescape View File

@ -53,21 +53,41 @@ class JointType(IntEnum):
params = {
# paths
'data_dir': '/data0/zhy/dataset/coco',
'vgg_path': '/data0/zhy/dataset/coco/vgg19-0-97_5004.ckpt',
'data_dir': './dataset',
'save_model_path': './checkpoints/',
'load_pretrain': False,
'pretrained_model_path': "",
# training params
'batch_size': 10,
'lr': 1e-4,
'lr_gamma': 0.1,
'lr_steps': '100000,200000,250000',
'lr_steps_NP': '250000',
# train type
'train_type': 'fix_loss_scale', # chose in ['clip_grad', 'fix_loss_scale']
'train_type_NP': 'clip_grad',
# vgg bn
'vgg_with_bn': False,
'vgg_path': './vgg_model/vgg19-0-97_5004.ckpt',
# if clip_grad
'GRADIENT_CLIP_TYPE': 1,
'GRADIENT_CLIP_VALUE': 10.0,
'loss_scale': 16386,
# optimizer and lr
'optimizer': "Adam", # chose in ['Momentum', 'Adam']
'optimizer_NP': "Momentum",
'group_params': True,
'group_params_NP': False,
'lr': 1e-4,
'lr_type': 'default', # chose in ["default", "cosine"]
'lr_gamma': 0.1, # if default
'lr_steps': '100000,200000,250000', # if default
'lr_steps_NP': '250000,300000', # if default
'warmup_epoch': 5, # if cosine
'max_epoch_train': 60,
'max_epoch_train_NP': 80,
'loss_scale': 16384,
# default param
'batch_size': 10,
'min_keypoints': 5,
'min_area': 32 * 32,
'insize': 368,
@ -75,9 +95,9 @@ params = {
'paf_sigma': 8,
'heatmap_sigma': 7,
'eva_num': 100,
'keep_checkpoint_max': 5,
'keep_checkpoint_max': 1,
'log_interval': 100,
'ckpt_interval': 663, # 5000,
'ckpt_interval': 5304,
'min_box_size': 64,
'max_box_size': 512,

128
model_zoo/official/cv/openpose/src/dataset.py
Unescape View File

@ -15,10 +15,10 @@
import os
import math
import random
import cv2
import numpy as np
import cv2
from pycocotools.coco import COCO as ReadJson
import mindspore.dataset as de
from src.config import JointType, params
@ -41,6 +41,7 @@ class txtdataset():
self.imgIds = random.sample(self.imgIds, n_samples)
print('{} images: {}'.format(mode, len(self)))
def __len__(self):
return len(self.imgIds)
@ -217,9 +218,9 @@ class txtdataset():
flipped_mask = cv2.flip(mask.astype(np.uint8), 1).astype('bool')
poses[:, :, 0] = img.shape[1] - 1 - poses[:, :, 0]
def swap_joints(poses, joint_type_1, joint_type_2):
tmp = poses[:, joint_type_1].copy()
poses[:, joint_type_1] = poses[:, joint_type_2]
def swap_joints(poses, joint_type_, joint_type_2):
tmp = poses[:, joint_type_].copy()
poses[:, joint_type_] = poses[:, joint_type_2]
poses[:, joint_type_2] = tmp
swap_joints(poses, JointType.LeftEye, JointType.RightEye)
@ -243,8 +244,10 @@ class txtdataset():
aug_img, ignore_mask, poses = self.flip_img(aug_img, ignore_mask, poses)
return aug_img, ignore_mask, poses
# ------------------------------------------------------------------
# ------------------------------- end -----------------------------------
# ------------------------------ Heatmap ------------------------------------
# return shape: (height, width)
def generate_gaussian_heatmap(self, shape, joint, sigma):
x, y = joint
@ -269,6 +272,38 @@ class txtdataset():
heatmaps = np.vstack((heatmaps, bg_heatmap[None]))
return heatmaps.astype('f')
def generate_gaussian_heatmap_fast(self, shape, joint, sigma):
x, y = joint
grid_x = np.tile(np.arange(shape[1]), (shape[0], 1))
grid_y = np.tile(np.arange(shape[0]), (shape[1], 1)).transpose()
grid_x = grid_x + 0.4375
grid_y = grid_y + 0.4375
grid_distance = (grid_x - x) ** 2 + (grid_y - y) ** 2
gaussian_heatmap = np.exp(-0.5 * grid_distance / sigma**2)
return gaussian_heatmap
def generate_heatmaps_fast(self, img, poses, heatmap_sigma):
resize_shape = (img.shape[0] // 8, img.shape[1] // 8)
heatmaps = np.zeros((0,) + resize_shape)
sum_heatmap = np.zeros(resize_shape)
for joint_index in range(len(JointType)):
heatmap = np.zeros(resize_shape)
for pose in poses:
if pose[joint_index, 2] > 0:
jointmap = self.generate_gaussian_heatmap_fast(resize_shape, pose[joint_index][:2]/8,
heatmap_sigma/8)
index_1 = jointmap > heatmap
heatmap[index_1] = jointmap[index_1]
index_2 = jointmap > sum_heatmap
sum_heatmap[index_2] = jointmap[index_2]
heatmaps = np.vstack((heatmaps, heatmap.reshape((1,) + heatmap.shape)))
bg_heatmap = 1 - sum_heatmap # background channel
heatmaps = np.vstack((heatmaps, bg_heatmap[None]))
return heatmaps.astype('f')
# ------------------------------ end ------------------------------------
# ------------------------------ PAF ------------------------------------
# return shape: (2, height, width)
def generate_constant_paf(self, shape, joint_from, joint_to, paf_width):
if np.array_equal(joint_from, joint_to): # same joint
@ -285,7 +320,7 @@ class txtdataset():
grid_y = np.tile(np.arange(shape[0]), (shape[1], 1)).transpose()
horizontal_inner_product = unit_vector[0] * (grid_x - joint_from[0]) + unit_vector[1] * (grid_y - joint_from[1])
horizontal_paf_flag = (horizontal_inner_product >= 0) & (horizontal_inner_product <= joint_distance)
vertical_inner_product = vertical_unit_vector[0] * (grid_x - joint_from[0]) + vertical_unit_vector[1] * \
vertical_inner_product = vertical_unit_vector[0] * (grid_x - joint_from[0]) + vertical_unit_vector[1] *\
(grid_y - joint_from[1])
vertical_paf_flag = np.abs(vertical_inner_product) <= paf_width # paf_width : 8
paf_flag = horizontal_paf_flag & vertical_paf_flag
@ -314,6 +349,55 @@ class txtdataset():
pafs = np.vstack((pafs, paf))
return pafs.astype('f')
def generate_constant_paf_fast(self, shape, joint_from, joint_to, paf_width):
if np.array_equal(joint_from, joint_to): # same joint
return np.zeros((2,) + shape[:-1])
joint_distance = np.linalg.norm(joint_to - joint_from)
unit_vector = (joint_to - joint_from) / joint_distance
rad = np.pi / 2
# [[0, 1], [-1, 0]]
rot_matrix = np.array([[np.cos(rad), np.sin(rad)], [-np.sin(rad), np.cos(rad)]])
# [[u_y], [-u_x]]
vertical_unit_vector = np.dot(rot_matrix, unit_vector)
grid_x = np.tile(np.arange(shape[1]), (shape[0], 1))
grid_y = np.tile(np.arange(shape[0]), (shape[1], 1)).transpose()
grid_x = grid_x + 0.4375
grid_y = grid_y + 0.4375
horizontal_inner_product = unit_vector[0] * (grid_x - joint_from[0]) + unit_vector[1] * (grid_y - joint_from[1])
horizontal_paf_flag = (horizontal_inner_product >= 0) & (horizontal_inner_product <= joint_distance)
vertical_inner_product = vertical_unit_vector[0] * (grid_x - joint_from[0]) + vertical_unit_vector[1] *\
(grid_y - joint_from[1])
vertical_paf_flag = np.abs(vertical_inner_product) <= paf_width # paf_width : 8/8 = 1
paf_flag = horizontal_paf_flag & vertical_paf_flag
constant_paf = np.stack((paf_flag, paf_flag)) *\
np.broadcast_to(unit_vector, shape[:-1] + (2,)).transpose(2, 0, 1)
return constant_paf
def generate_pafs_fast(self, img, poses, paf_sigma):
resize_shape = (img.shape[0]//8, img.shape[1]//8, 3)
pafs = np.zeros((0,) + resize_shape[:-1])
for limb in params['limbs_point']:
paf = np.zeros((2,) + resize_shape[:-1])
paf_flags = np.zeros(paf.shape) # for constant paf
for pose in poses:
joint_from, joint_to = pose[limb]
if joint_from[2] > 0 and joint_to[2] > 0:
limb_paf = self.generate_constant_paf_fast(resize_shape, joint_from[:2]/8, joint_to[:2]/8, paf_sigma/8) # [2, 368, 368]
limb_paf_flags = limb_paf != 0
paf_flags += np.broadcast_to(limb_paf_flags[0] | limb_paf_flags[1], limb_paf.shape)
paf += limb_paf
index_1 = paf_flags > 0
paf[index_1] /= paf_flags[index_1]
pafs = np.vstack((pafs, paf))
return pafs.astype('f')
# ------------------------------ end ------------------------------------
def get_img_annotation(self, ind=None, img_id=None):
annotations = None
@ -389,14 +473,18 @@ class txtdataset():
resized_img, ignore_mask, resized_poses = self.resize_data(img, ignore_mask, poses,
shape=(self.insize, self.insize))
heatmaps = self.generate_heatmaps(resized_img, resized_poses, params['heatmap_sigma'])
pafs = self.generate_pafs(resized_img, resized_poses, params['paf_sigma']) # params['paf_sigma']: 8
# heatmaps = self.generate_heatmaps(resized_img, resized_poses, params['heatmap_sigma'])
# resized_heatmaps = self.resize_output(heatmaps)
resized_heatmaps = self.generate_heatmaps_fast(resized_img, resized_poses, params['heatmap_sigma'])
# pafs = self.generate_pafs(resized_img, resized_poses, params['paf_sigma'])
# resized_pafs = self.resize_output(pafs)
resized_pafs = self.generate_pafs_fast(resized_img, resized_poses, params['paf_sigma'])
ignore_mask = cv2.morphologyEx(ignore_mask.astype('uint8'), cv2.MORPH_DILATE, np.ones((16, 16))).astype('bool')
resized_pafs = self.resize_output(pafs)
resized_heatmaps = self.resize_output(heatmaps)
resized_ignore_mask = self.resize_output(ignore_mask)
return resized_img, resized_pafs, resized_heatmaps, resized_ignore_mask
def preprocess(self, img):
@ -459,7 +547,6 @@ class DistributedSampler():
def __len__(self):
return self.num_samplers
def valdata(jsonpath, imgpath, rank, group_size, mode='val', maskpath=''):
#cv2.setNumThreads(0)
val = ReadJson(jsonpath)
@ -470,23 +557,6 @@ def valdata(jsonpath, imgpath, rank, group_size, mode='val', maskpath=''):
return ds
def openpose(jsonpath, imgpath, maskpath, per_batch_size, max_epoch, rank, group_size, mode='train'):
train = ReadJson(jsonpath)
num_parallel = 48
if group_size > 1:
num_parallel = 20
dataset = txtdataset(train, imgpath, maskpath, params['insize'], mode=mode)
sampler = DistributedSampler(dataset, rank, group_size)
de_dataset = de.GeneratorDataset(dataset, ["image", "pafs", "heatmaps", "ignore_mask"],
num_parallel_workers=num_parallel, sampler=sampler, shuffle=True)
de_dataset = de_dataset.project(columns=["image", "pafs", "heatmaps", "ignore_mask"])
de_dataset = de_dataset.batch(batch_size=per_batch_size, drop_remainder=True, num_parallel_workers=num_parallel)
steap_pre_epoch = de_dataset.get_dataset_size()
de_dataset = de_dataset.repeat(max_epoch)
return de_dataset, steap_pre_epoch
def create_dataset(jsonpath, imgpath, maskpath, batch_size, rank, group_size, mode='train', repeat_num=1, shuffle=True,
multiprocessing=True, num_worker=20):

81
model_zoo/official/cv/openpose/src/gen_ignore_mask.py
Unescape View File

@ -15,7 +15,6 @@
import os
import argparse
import cv2
import numpy as np
from tqdm import tqdm
from pycocotools.coco import COCO as ReadJson
@ -23,44 +22,44 @@ from pycocotools.coco import COCO as ReadJson
from config import params
class DataLoader():
def __init__(self, coco, dir_name, data_mode='train'):
self.train = coco
def __init__(self, train_, dir_name, mode_='train'):
self.train = train_
self.dir_name = dir_name
assert data_mode in ['train', 'val'], 'Data loading mode is invalid.'
self.mode = data_mode
self.catIds = coco.getCatIds() # catNms=['person']
self.imgIds = sorted(coco.getImgIds(catIds=self.catIds))
assert mode_ in ['train', 'val'], 'Data loading mode is invalid.'
self.mode = mode_
self.catIds = train_.getCatIds() # catNms=['person']
self.imgIds = sorted(train_.getImgIds(catIds=self.catIds))
def __len__(self):
return len(self.imgIds)
def gen_masks(self, image, anns):
_mask_all = np.zeros(image.shape[:2], 'bool')
_mask_miss = np.zeros(image.shape[:2], 'bool')
for ann in anns:
def gen_masks(self, image_, annotations_):
mask_all_1 = np.zeros(image_.shape[:2], 'bool')
mask_miss_1 = np.zeros(image_.shape[:2], 'bool')
for ann in annotations_:
mask = self.train.annToMask(ann).astype('bool')
if ann['iscrowd'] == 1:
intxn = _mask_all & mask
_mask_miss = np.bitwise_or(_mask_miss.astype(int), np.subtract(mask, intxn, dtype=np.int32))
_mask_all = np.bitwise_or(_mask_all.astype(int), mask.astype(int))
intxn = mask_all_1 & mask
mask_miss_1 = np.bitwise_or(mask_miss_1.astype(int), np.subtract(mask, intxn, dtype=np.int32))
mask_all_1 = np.bitwise_or(mask_all_1.astype(int), mask.astype(int))
elif ann['num_keypoints'] < params['min_keypoints'] or ann['area'] <= params['min_area']:
_mask_all = np.bitwise_or(_mask_all.astype(int), mask.astype(int))
_mask_miss = np.bitwise_or(_mask_miss.astype(int), mask.astype(int))
mask_all_1 = np.bitwise_or(mask_all_1.astype(int), mask.astype(int))
mask_miss_1 = np.bitwise_or(mask_miss_1.astype(int), mask.astype(int))
else:
_mask_all = np.bitwise_or(_mask_all.astype(int), mask.astype(int))
return _mask_all, _mask_miss
mask_all_1 = np.bitwise_or(mask_all_1.astype(int), mask.astype(int))
return mask_all_1, mask_miss_1
def dwaw_gen_masks(self, image, mask, color=(0, 0, 1)):
def dwaw_gen_masks(self, image_, mask, color=(0, 0, 1)):
bimsk = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
mskd = image * bimsk.astype(np.int32)
mskd = image_ * bimsk.astype(np.int32)
clmsk = np.ones(bimsk.shape) * bimsk
for ind in range(3):
clmsk[:, :, ind] = clmsk[:, :, ind] * color[ind] * 255
image = image + 0.7 * clmsk - 0.7 * mskd
return image.astype(np.uint8)
for index_1 in range(3):
clmsk[:, :, index_1] = clmsk[:, :, index_1] * color[index_1] * 255
image_ = image_ + 0.7 * clmsk - 0.7 * mskd
return image_.astype(np.uint8)
def draw_masks_and_keypoints(self, image, anns):
for ann in anns:
def draw_masks_and_keypoints(self, image_, annotations_):
for ann in annotations_:
# masks
mask = self.train.annToMask(ann).astype(np.uint8)
if ann['iscrowd'] == 1:
@ -70,30 +69,30 @@ class DataLoader():
else:
color = (1, 0, 0)
bimsk = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
mskd = image * bimsk.astype(np.int32)
mskd = image_ * bimsk.astype(np.int32)
clmsk = np.ones(bimsk.shape) * bimsk
for ind in range(3):
clmsk[:, :, ind] = clmsk[:, :, ind] * color[ind] * 255
image = image + 0.7 * clmsk - 0.7 * mskd
for index_1 in range(3):
clmsk[:, :, index_1] = clmsk[:, :, index_1] * color[index_1] * 255
image_ = image_ + 0.7 * clmsk - 0.7 * mskd
# keypoints
for x, y, v in np.array(ann['keypoints']).reshape(-1, 3):
if v == 1:
cv2.circle(image, (x, y), 3, (255, 255, 0), -1)
cv2.circle(image_, (x, y), 3, (255, 255, 0), -1)
elif v == 2:
cv2.circle(image, (x, y), 3, (255, 0, 255), -1)
return image.astype(np.uint8)
cv2.circle(image_, (x, y), 3, (255, 0, 255), -1)
return image_.astype(np.uint8)
def get_img_annotation(self, ind=None, image_id=None):
def get_img_annotation(self, ind=None, img_id_=None):
if ind is not None:
image_id = self.imgIds[ind]
img_id_ = self.imgIds[ind]
anno_ids = self.train.getAnnIds(imgIds=[image_id])
_annotations = self.train.loadAnns(anno_ids)
anno_ids = self.train.getAnnIds(imgIds=[img_id_])
annotations_ = self.train.loadAnns(anno_ids)
img_file = os.path.join(params['data_dir'], self.dir_name, self.train.loadImgs([image_id])[0]['file_name'])
_image = cv2.imread(img_file)
return _image, _annotations, image_id
img_file = os.path.join(params['data_dir'], self.dir_name, self.train.loadImgs([img_id_])[0]['file_name'])
image_ = cv2.imread(img_file)
return image_, annotations_, img_id_
if __name__ == '__main__':
@ -107,7 +106,7 @@ if __name__ == '__main__':
path_list = [args.train_ann, args.val_ann, args.train_dir, args.val_dir]
for index, mode in enumerate(['train', 'val']):
train = ReadJson(path_list[index])
data_loader = DataLoader(train, path_list[index+2], mode=mode)
data_loader = DataLoader(train, path_list[index+2], mode_=mode)
save_dir = os.path.join(params['data_dir'], 'ignore_mask_{}'.format(mode))
if not os.path.exists(save_dir):

164
model_zoo/official/cv/openpose/src/loss.py
Unescape View File

@ -12,37 +12,53 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
import time
import mindspore.nn as nn
from mindspore.ops import operations as P
from mindspore.nn.loss.loss import _Loss
from mindspore.train.callback import Callback
from mindspore.ops import functional as F
from mindspore.ops import composite as C
from mindspore.communication.management import get_group_size
from mindspore.context import ParallelMode
from mindspore import context
from mindspore.parallel._utils import _get_device_num, _get_parallel_mode, _get_gradients_mean
from mindspore.nn.wrap.grad_reducer import DistributedGradReducer
from src.config import params
context.set_context(mode=context.GRAPH_MODE, save_graphs=True)
time_stamp_init = False
time_stamp_first = 0
grad_scale = C.MultitypeFuncGraph("grad_scale")
_grad_overflow = C.MultitypeFuncGraph("_grad_overflow")
reciprocal = P.Reciprocal()
GRADIENT_CLIP_TYPE = 1
GRADIENT_CLIP_VALUE = 1.0
@grad_scale.register("Tensor", "Tensor")
def tensor_grad_scale(scale, grad):
return grad * F.cast(reciprocal(scale), F.dtype(grad))
@grad_scale.register("Tensor", "RowTensor")
def tensor_grad_scale_row_tensor(scale, grad):
return RowTensor(grad.indices,
grad.values * F.cast(reciprocal(scale), F.dtype(grad.values)),
grad.dense_shape)
GRADIENT_CLIP_TYPE = params['GRADIENT_CLIP_TYPE']
GRADIENT_CLIP_VALUE = params['GRADIENT_CLIP_VALUE']
clip_grad = C.MultitypeFuncGraph("clip_grad")
@clip_grad.register("Number", "Number", "Tensor")
def _clip_grad(clip_type, clip_value, grad):
"""
Clip gradients.
Inputs:
clip_type (int): The way to clip, 0 for 'value', 1 for 'norm'.
clip_value (float): Specifies how much to clip.
grad (tuple[Tensor]): Gradients.
Outputs:
tuple[Tensor]: clipped gradients.
"""
if clip_type not in (0, 1):
return grad
dt = F.dtype(grad)
if clip_type == 0:
new_grad = C.clip_by_value(grad, F.cast(F.tuple_to_array((-clip_value,)), dt),
F.cast(F.tuple_to_array((clip_value,)), dt))
else:
new_grad = nn.ClipByNorm()(grad, F.cast(F.tuple_to_array((clip_value,)), dt))
return new_grad
class openpose_loss(_Loss):
def __init__(self):
super(openpose_loss, self).__init__()
@ -99,109 +115,49 @@ class openpose_loss(_Loss):
return total_loss, heatmaps_loss, pafs_loss
class Depend_network(nn.Cell):
def __init__(self, network):
super(Depend_network, self).__init__()
class BuildTrainNetwork(nn.Cell):
def __init__(self, network, criterion):
super(BuildTrainNetwork, self).__init__()
self.network = network
self.criterion = criterion
def construct(self, *args):
loss, _, _ = self.network(*args) # loss, heatmaps_loss, pafs_loss
def construct(self, input_data, gt_paf, gt_heatmap, mask):
logit_pafs, logit_heatmap = self.network(input_data)
loss, _, _ = self.criterion(logit_pafs, logit_heatmap, gt_paf, gt_heatmap, mask)
return loss
#loss = self.criterion(logit_pafs, logit_heatmap, gt_paf, gt_heatmap, mask)
# return loss, heatmaps_loss, pafs_loss
class TrainingWrapper(nn.Cell):
def __init__(self, network, optimizer, sens=1):
super(TrainingWrapper, self).__init__(auto_prefix=False)
class TrainOneStepWithClipGradientCell(nn.Cell):
'''TrainOneStepWithClipGradientCell'''
def __init__(self, network, optimizer, sens=1.0):
super(TrainOneStepWithClipGradientCell, self).__init__(auto_prefix=False)
self.network = network
self.depend_network = Depend_network(network)
# self.weights = ms.ParameterTuple(network.trainable_params())
self.network.set_grad()
self.network.add_flags(defer_inline=True)
self.weights = optimizer.parameters
self.optimizer = optimizer
self.grad = C.GradOperation(get_by_list=True, sens_param=True)
self.hyper_map = C.HyperMap()
self.sens = sens
self.reducer_flag = False
self.grad_reducer = None
self.print = P.Print()
self.parallel_mode = context.get_auto_parallel_context("parallel_mode")
if self.parallel_mode in [ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL]:
parallel_mode = _get_parallel_mode()
if parallel_mode in (ParallelMode.DATA_PARALLEL, ParallelMode.HYBRID_PARALLEL):
self.reducer_flag = True
if self.reducer_flag:
mean = context.get_auto_parallel_context("gradients_mean")
#if mean.get_device_num_is_set():
# if mean:
#degree = context.get_auto_parallel_context("device_num")
# else:
degree = get_group_size()
self.grad_reducer = nn.DistributedGradReducer(optimizer.parameters, mean, degree)
def construct(self, *args):
mean = _get_gradients_mean()
degree = _get_device_num()
self.grad_reducer = DistributedGradReducer(optimizer.parameters, mean, degree)
def construct(self, *inputs):
weights = self.weights
loss, heatmaps_loss, pafs_loss = self.network(*args)
loss = self.network(*inputs)
sens = P.Fill()(P.DType()(loss), P.Shape()(loss), self.sens)
#grads = self.grad(self.network, weights)(*args, sens)
grads = self.grad(self.depend_network, weights)(*args, sens)
grads = self.grad(self.network, weights)(*inputs, sens)
grads = self.hyper_map(F.partial(clip_grad, GRADIENT_CLIP_TYPE, GRADIENT_CLIP_VALUE), grads)
if self.reducer_flag:
# apply grad reducer on grads
grads = self.grad_reducer(grads)
#return F.depend(loss, self.optimizer(grads))
# for grad in grads:
# self.print(grad)
loss = F.depend(loss, self.optimizer(grads))
return loss, heatmaps_loss, pafs_loss
class BuildTrainNetwork(nn.Cell):
def __init__(self, network, criterion):
super(BuildTrainNetwork, self).__init__()
self.network = network
self.criterion = criterion
def construct(self, input_data, gt_paf, gt_heatmap, mask):
logit_pafs, logit_heatmap = self.network(input_data)
loss, _, _ = self.criterion(logit_pafs, logit_heatmap, gt_paf, gt_heatmap, mask)
return loss
class LossCallBack(Callback):
"""
Monitor the loss in training.
If the loss is NAN or INF terminating training.
Note:
If per_print_times is 0 do not print loss.
Args:
per_print_times (int): Print loss every times. Default: 1.
"""
def __init__(self, per_print_times=1):
super(LossCallBack, self).__init__()
if not isinstance(per_print_times, int) or per_print_times < 0:
raise ValueError("print_step must be int and >= 0.")
self._per_print_times = per_print_times
self.count = 0
self.loss_sum = 0
global time_stamp_init, time_stamp_first
if not time_stamp_init:
time_stamp_first = time.time()
time_stamp_init = True
def step_end(self, run_context):
cb_params = run_context.original_args()
loss = cb_params.net_outputs.asnumpy()
self.count += 1
self.loss_sum += float(loss)
cur_step_in_epoch = (cb_params.cur_step_num - 1) % cb_params.batch_num + 1
if self.count >= 1:
global time_stamp_first
time_stamp_current = time.time()
loss = self.loss_sum/self.count
loss_file = open("./loss.log", "a+")
loss_file.write("%lu epoch: %s step: %s ,loss: %.5f" %
(time_stamp_current - time_stamp_first, cb_params.cur_epoch_num, cur_step_in_epoch,
loss))
loss_file.write("\n")
loss_file.close()
self.count = 0
self.loss_sum = 0
return F.depend(loss, self.optimizer(grads))

24
model_zoo/official/cv/openpose/src/openposenet.py
Unescape View File

@ -17,19 +17,18 @@ from mindspore.nn import Conv2d, ReLU
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from mindspore.ops import operations as P
from mindspore import context
context.set_context(mode=context.GRAPH_MODE, save_graphs=True)
#selfCat = P.Concat(axis=1)
time_stamp_init = False
time_stamp_first = 0
loadvgg = 1
class OpenPoseNet(nn.Cell):
insize = 368
def __init__(self, vggpath=''):
def __init__(self, vggpath='', vgg_with_bn=False):
super(OpenPoseNet, self).__init__()
self.base = Base_model()
self.base = Base_model(vgg_with_bn=vgg_with_bn)
self.stage_1 = Stage_1()
self.stage_2 = Stage_x()
self.stage_3 = Stage_x()
@ -39,23 +38,15 @@ class OpenPoseNet(nn.Cell):
self.shape = P.Shape()
self.cat = P.Concat(axis=1)
self.print = P.Print()
# for m in self.modules():
# if isinstance(m, Conv2d):
# init.constant_(m.bias, 0)
if loadvgg and vggpath:
param_dict = load_checkpoint(vggpath)
param_dict_new = {}
trans_name = 'base.vgg_base.'
for key, values in param_dict.items():
#print('key:',key,self.shape(values))
if key.startswith('moments.'):
continue
elif key.startswith('network.'):
param_dict_new[trans_name+key[17:]] = values
# else:
# param_dict_new[key] = values
#print(param_dict_new)
load_param_into_net(self.base.vgg_base, param_dict_new)
def construct(self, x):
@ -205,20 +196,17 @@ class VGG_Base_MS(nn.Cell):
return x
class Base_model(nn.Cell):
def __init__(self):
def __init__(self, vgg_with_bn=False):
super(Base_model, self).__init__()
#cfgs_zh = {'16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512]}
cfgs_zh = {'19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512]}
#cfgs_zh = {'16': [64, 64,128, 128, 256, 256, 256, 512, 512, 512]}
self.vgg_base = Vgg(cfgs_zh['19'], batch_norm=False)
#self.vgg_base = VGG_Base()
#self.vgg_base = VGG_Base_MS()
self.vgg_base = Vgg(cfgs_zh['19'], batch_norm=vgg_with_bn)
self.conv4_3_CPM = Conv2d(in_channels=512, out_channels=256, kernel_size=3, stride=1, pad_mode='same',
has_bias=True)
self.conv4_4_CPM = Conv2d(in_channels=256, out_channels=128, kernel_size=3, stride=1, pad_mode='same',
has_bias=True)
self.relu = ReLU()
def construct(self, x):
x = self.vgg_base(x)
x = self.relu(self.conv4_3_CPM(x))

124
model_zoo/official/cv/openpose/src/utils.py
Unescape View File

@ -1,27 +1,11 @@
# Copyright 2020 Huawei Technologies Co., Ltd
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
import argparse
import os
import math
import time
import numpy as np
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from mindspore.train.callback import LossMonitor
from mindspore.train.callback import LossMonitor, Callback
from mindspore.common.tensor import Tensor
from src.config import params
from mindspore.common import dtype as mstype
class MyLossMonitor(LossMonitor):
def __init__(self, per_print_times=1):
@ -32,6 +16,7 @@ class MyLossMonitor(LossMonitor):
def step_end(self, run_context):
cb_params = run_context.original_args()
loss = cb_params.net_outputs
if isinstance(loss, (tuple, list)):
@ -47,63 +32,76 @@ class MyLossMonitor(LossMonitor):
raise ValueError("epoch: {} step: {}. Invalid loss, terminating training.".format(
cb_params.cur_epoch_num, cur_step_in_epoch))
if self._per_print_times != 0 and cb_params.cur_step_num % self._per_print_times == 0:
# print("epoch: %s step: %s, loss is %s, step time: %.3f s." % (cb_params.cur_epoch_num, cur_step_in_epoch,
# loss,
# (time.time() - self._start_time)), flush=True)
self._loss_list.append(loss)
if cb_params.cur_step_num % 100 == 0:
print("epoch: %s, steps: [%s] mean loss is: %s"%(cb_params.cur_epoch_num, cur_step_in_epoch,
np.array(self._loss_list).mean()), flush=True)
print("epoch: %s, steps: [%s], mean loss is: %s"%(cb_params.cur_epoch_num, cur_step_in_epoch,
np.array(self._loss_list).mean()), flush=True)
self._loss_list = []
self._start_time = time.time()
class MyScaleSensCallback(Callback):
'''MyLossScaleCallback'''
def __init__(self, loss_scale_list, epoch_list):
super(MyScaleSensCallback, self).__init__()
self.loss_scale_list = loss_scale_list
self.epoch_list = epoch_list
self.scaling_sens = loss_scale_list[0]
def parse_args():
"""Parse train arguments."""
parser = argparse.ArgumentParser('mindspore openpose training')
# dataset related
parser.add_argument('--train_dir', type=str, default='train2017', help='train data dir')
parser.add_argument('--train_ann', type=str, default='person_keypoints_train2017.json',
help='train annotations json')
parser.add_argument('--group_size', type=int, default=1, help='world size of distributed')
args, _ = parser.parse_known_args()
args.jsonpath_train = os.path.join(params['data_dir'], 'annotations/' + args.train_ann)
args.imgpath_train = os.path.join(params['data_dir'], args.train_dir)
args.maskpath_train = os.path.join(params['data_dir'], 'ignore_mask_train')
return args
def get_lr(lr, lr_gamma, steps_per_epoch, max_epoch_train, lr_steps, group_size):
lr_stage = np.array([lr] * steps_per_epoch * max_epoch_train).astype('f')
for step in lr_steps:
step //= group_size
lr_stage[step:] *= lr_gamma
def epoch_end(self, run_context):
cb_params = run_context.original_args()
epoch = cb_params.cur_epoch_num
for i, _ in enumerate(self.epoch_list):
if epoch >= self.epoch_list[i]:
self.scaling_sens = self.loss_scale_list[i+1]
else:
break
scaling_sens_tensor = Tensor(self.scaling_sens, dtype=mstype.float32)
cb_params.train_network.set_sense_scale(scaling_sens_tensor)
print("Epoch: set train network scale sens to {}".format(self.scaling_sens))
def _linear_warmup_learning_rate(current_step, warmup_steps, base_lr, init_lr):
lr_inc = (float(base_lr) - float(init_lr)) / float(warmup_steps)
learning_rate = float(init_lr) + lr_inc * current_step
return learning_rate
def _a_cosine_learning_rate(current_step, base_lr, warmup_steps, decay_steps):
base = float(current_step - warmup_steps) / float(decay_steps)
learning_rate = (1 + math.cos(base * math.pi)) / 2 * base_lr
return learning_rate
def _dynamic_lr(base_lr, total_steps, warmup_steps, warmup_ratio=1 / 3):
lr = []
for i in range(total_steps):
if i < warmup_steps:
lr.append(_linear_warmup_learning_rate(i, warmup_steps, base_lr, base_lr * warmup_ratio))
else:
lr.append(_a_cosine_learning_rate(i, base_lr, warmup_steps, total_steps))
return lr
def get_lr(lr, lr_gamma, steps_per_epoch, max_epoch_train, lr_steps, group_size, lr_type='default', warmup_epoch=5):
if lr_type == 'default':
lr_stage = np.array([lr] * steps_per_epoch * max_epoch_train).astype('f')
for step in lr_steps:
step //= group_size
lr_stage[step:] *= lr_gamma
elif lr_type == 'cosine':
lr_stage = _dynamic_lr(lr, steps_per_epoch * max_epoch_train, warmup_epoch * steps_per_epoch,
warmup_ratio=1 / 3)
lr_stage = np.array(lr_stage).astype('f')
else:
raise ValueError("lr type {} is not support.".format(lr_type))
lr_base = lr_stage.copy()
lr_base = lr_base / 4
lr_vgg = lr_base.copy()
vgg_freeze_step = 2000
vgg_freeze_step = 2000 // group_size
lr_vgg[:vgg_freeze_step] = 0
return lr_stage, lr_base, lr_vgg
# zhang add
def adjust_learning_rate(init_lr, lr_gamma, steps_per_epoch, max_epoch_train, stepvalues):
lr_stage = np.array([init_lr] * steps_per_epoch * max_epoch_train).astype('f')
for epoch in stepvalues:
lr_stage[epoch * steps_per_epoch:] *= lr_gamma
lr_base = lr_stage.copy()
lr_base = lr_base / 4
lr_vgg = lr_base.copy()
vgg_freeze_step = 2000
lr_vgg[:vgg_freeze_step] = 0
return lr_stage, lr_base, lr_vgg

104
model_zoo/official/cv/openpose/train.py
Unescape View File

@ -13,26 +13,38 @@
# limitations under the License.
# ============================================================================
import os
import argparse
import mindspore
from mindspore import context
from mindspore.context import ParallelMode
from mindspore.communication.management import init, get_rank, get_group_size
from mindspore.train import Model
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, TimeMonitor
from mindspore.nn.optim import Adam, Momentum
from mindspore.train.loss_scale_manager import FixedLossScaleManager
from mindspore.nn.optim import Adam
from src.dataset import create_dataset
from src.openposenet import OpenPoseNet
from src.loss import openpose_loss, BuildTrainNetwork
from src.loss import openpose_loss, BuildTrainNetwork, TrainOneStepWithClipGradientCell
from src.config import params
from src.utils import parse_args, get_lr, load_model, MyLossMonitor
from src.utils import get_lr, load_model, MyLossMonitor
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", save_graphs=False)
parser = argparse.ArgumentParser('mindspore openpose training')
parser.add_argument('--train_dir', type=str, default='train2017', help='train data dir')
parser.add_argument('--train_ann', type=str, default='person_keypoints_train2017.json',
help='train annotations json')
parser.add_argument('--group_size', type=int, default=1, help='world size of distributed')
args, _ = parser.parse_known_args()
args.jsonpath_train = os.path.join(params['data_dir'], 'annotations/' + args.train_ann)
args.imgpath_train = os.path.join(params['data_dir'], args.train_dir)
args.maskpath_train = os.path.join(params['data_dir'], 'ignore_mask_train')
def train():
"""Train function."""
args = parse_args()
args.outputs_dir = params['save_model_path']
@ -46,11 +58,15 @@ def train():
args.outputs_dir = os.path.join(args.outputs_dir, "ckpt_0/")
args.rank = 0
# with out loss_scale
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(',')))
@ -58,9 +74,7 @@ def train():
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)
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"])
@ -72,7 +86,7 @@ def 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,
@ -90,35 +104,63 @@ def train():
lr_stage, lr_base, lr_vgg = get_lr(params['lr'] * args.group_size,
params['lr_gamma'],
steps_per_epoch,
params["max_epoch_train"],
args.max_epoch,
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'])
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())
callback_list = [MyLossMonitor(), time_cb, ckpoint_cb]
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(params["max_epoch_train"], de_dataset_train, callbacks=callback_list,
model.train(args.max_epoch, de_dataset_train, callbacks=callback_list,
dataset_sink_mode=False)
return 0
if __name__ == "__main__":
# mindspore.common.seed.set_seed(1)
mindspore.common.seed.set_seed(1)
train()

Write Preview
Loading…
Cancel
Save