# cfg.MODEL.scale_list,
# cfg.LOSS.SCORE,
# cfg.LOSS.PAIR,
# )
#
# device = torch.device("cuda:0")
# model = det.to(device)
# checkpoint = torch.load(resume)
# det.load_state_dict(checkpoint["state_dict"])
#
# img = torch.from_numpy(img.transpose((2, 0, 1)))[None, :].to(
# device, dtype=torch.float
# )
# score_maps, orint_maps = det(img)
#
# traced_script_module = torch.jit.trace(det, img)
# traced_script_module.save("det.pt")
resume = "/home/wang/workspace/RFNET/runs/07_25_14_05/model/e001_NN_0.283_NNT_0.337_NNDR_0.692_MeanMS_0.437_des.pth.tar"
des = HardNetNeiMask(cfg.HARDNET.MARGIN, cfg.MODEL.COO_THRSH)
device = torch.device("cuda")
des = des.to(device).eval()
checkpoint = torch.load(resume)
des.load_state_dict(checkpoint["state_dict"])
device = torch.device("cuda:0")
im_patches = torch.randn((512, 1, 32, 32)).to(device)
im_des = des(im_patches)
traced_script_module2 = torch.jit.trace(des, im_patches)
traced_script_module2.save("des.pt")
cfg.MODEL.GAUSSIAN_KSIZE,
cfg.MODEL.GAUSSIAN_SIGMA,
cfg.MODEL.KSIZE,
cfg.MODEL.padding,
cfg.MODEL.dilation,
cfg.MODEL.scale_list,
cfg.LOSS.SCORE,
cfg.LOSS.PAIR,
)
des = HardNetNeiMask(cfg.HARDNET.MARGIN, cfg.MODEL.COO_THRSH)
if mgpu:
det = torch.nn.DataParallel(det)
des = torch.nn.DataParallel(des)
det = det.to(device=device)
des = des.to(device=device)
###############################################################################
# Load train data
###############################################################################
PPT = [cfg.PROJ.TRAIN_PPT, (cfg.PROJ.TRAIN_PPT + cfg.PROJ.EVAL_PPT)]
print(f"{gct()} : Loading traning data")
train_data = DataLoader(
HpatchDataset(
data_type="train",
PPT=PPT,
use_all=cfg.PROJ.TRAIN_ALL,
csv_file=cfg[cfg.PROJ.TRAIN]["csv"],
root_dir=cfg[cfg.PROJ.TRAIN]["root"],
transform=transforms.Compose([