def __init__(self, params):
assert isinstance(params, Params)
self.params = params
if self.params.backbone == 'resnet101':
model = poseNet(101)
elif self.params.backbone == 'resnet50':
model = poseNet(50)
for name, module in model.named_children():
for para in module.parameters():
para.requires_grad = False
# load model
self.model = model
ckpt = self.params.ckpt
if ckpt is not None:
self._load_ckpt(ckpt)
logger.info('Load ckpt from {}'.format(ckpt))
self.model = nn.DataParallel(self.model, device_ids=self.params.gpus)
self.model = self.model.cuda(device=self.params.gpus[0])
self.model.eval()
self.model.module.freeze_bn()
sys.path.append(root_path)
from network.posenet import poseNet
from evaluate.tester import Tester
backbone = 'resnet101'
# Set Training parameters
params = Tester.TestParams()
params.subnet_name = 'both'
params.inp_size = 480 # input picture size = (inp_size, inp_size)
params.coeff = 2
params.in_thres = 0.21
params.coco_root = '/data/COCO/'
params.testresult_write_json = False # Whether to write json result
params.coco_result_filename = './demo/multipose_coco2017_results.json'
params.ckpt = './demo/models/ckpt_baseline_resnet101.h5'
# model
if backbone == 'resnet101':
model = poseNet(101)
elif backbone == 'resnet50':
model = poseNet(50)
for name, module in model.named_children():
for para in module.parameters():
para.requires_grad = False
tester = Tester(model, params)
tester.coco_eval() # pic_test
params.save_dir = './extra/models/{}'.format(params.exp_name)
params.ckpt = './demo/models/ckpt_baseline_resnet101.h5'
params.max_epoch = 40
params.init_lr = 1.0e-3
params.lr_decay = 0.9
params.gpus = [0]
params.batch_size = 8 * len(params.gpus)
params.val_nbatch_end_epoch = 2000
params.print_freq = 1000
# model
if backbone == 'resnet101':
model = poseNet(101, prn_node_count=node_count, prn_coeff=coeff)
elif backbone == 'resnet50':
model = poseNet(50, prn_node_count=node_count, prn_coeff=coeff)
# Train Key-point Subnet, Fix the weights in detection subnet (RetinaNet)
for name, module in model.named_children():
if name not in prn_para:
for para in module.parameters():
para.requires_grad = False
print("Loading dataset...")
# load training data
coco_train = COCO(
os.path.join(coco_root, 'annotations/person_keypoints_train2017.json'))
train_data = DataLoader(dataset=PRN_CocoDataset(
coco_train,
