def train(self):
lr_scheduer = {1: 1e-3, 2: 2e-3, 3: 1e-3, 60: 1e-4, 120: 1e-5}
# train
self.model.train()
for epoch in range(self.epochs):
if epoch in lr_scheduer:
self.set_lr(lr_scheduer[epoch])
self.train_epoch(epoch)
file = f"{jobdir}/models/{epoch + 1}.pth"
common.mkdirs_from_file_path(file)
torch.save(self.model.module.state_dict(), file)
# load image and forward
image = common.imread(file)
objs = eval_tool.detect_image(model, image, mean, std, 0.01)
# summary to all_result_dict
image_pred = []
for obj in objs:
image_pred.append(obj.xywh + [obj.score])
# build all_result_dict
if key not in all_result_dict:
all_result_dict[key] = {}
all_result_dict[key][file_name] = np.array(image_pred)
log.info("{} / {}".format(i + 1, total_file))
# write matlab format
path = f"{jobdir}/matlab_eval_{trial_name}/{key}/{file_name}.txt"
common.mkdirs_from_file_path(path)
with open(path, "w") as f:
f.write(f"/{key}/{file_name}\n{len(image_pred)}\n")
for item in image_pred:
f.write("{} {} {} {} {}\n".format(*item))
# eval map of IoU0.5
aps = evaluation.eval_map(all_result_dict, all=False)
log.info("\n" "Easy: {}\n" "Medium: {}\n" "Hard: {}".format(*aps))
class OnnxModel(nn.Module):
def __init__(self, **kwargs):
super(OnnxModel, self).__init__()
self.model = DBFace(**kwargs)
self.model.load(f"{jobdir}/models/150.pth")
def forward(self, x):
center, box, landmark = self.model(x)
center_sigmoid = torch.sigmoid(center)
center_maxpool = F.max_pool2d(center_sigmoid,
kernel_size=3,
padding=1,
stride=1)
box = torch.exp(box)
return center_sigmoid, center_maxpool, box, landmark
model = OnnxModel(has_landmark=True, wide=64, has_ext=True, upmode="UCBA")
model.eval()
model.cuda()
common.mkdirs_from_file_path(f"{jobdir}/model.onnx")
dummy = torch.zeros((1, 3, 32, 32)).cuda()
torch.onnx.export(model,
dummy,
f"{jobdir}/model.onnx",
output_names=["hm", "pool_hm", "tlrb", "landmark"],
opset_version=11)