def criterion(tgt_vocab_size, use_cuda,loss):
weight = torch.ones(tgt_vocab_size)
weight[dict.PAD] = 0
if loss=='focal_loss':
crit = focal_loss.FocalLoss(gamma=2)
else:
crit = nn.CrossEntropyLoss(weight, size_average=False)
if use_cuda:
crit.cuda()
return crit
def __init__(self, opt: Options, writer):
self.opt = opt
self.writer = writer
self.global_step = 0
self.detector = networks_united.KeypointDetector(self.opt).to(self.opt.device)
# self.coarse_ce_criteria = nn.CrossEntropyLoss()
self.coarse_ce_criteria = focal_loss.FocalLoss(alpha=0.5, gamma=2, reduction='mean')
# self.fine_ce_criteria = nn.CrossEntropyLoss()
# multi-gpu training
if len(opt.gpu_ids) > 1:
self.detector = nn.DataParallel(self.detector, device_ids=opt.gpu_ids)
# self.coarse_ce_criteria = nn.DataParallel(self.coarse_ce_criteria, device_ids=opt.gpu_ids)
# self.fine_ce_criteria = nn.DataParallel(self.fine_ce_criteria, device_ids=opt.gpu_ids)
# learning rate_control
self.old_lr_detector = self.opt.lr
self.optimizer_detector = torch.optim.Adam(self.detector.parameters(),
lr=self.old_lr_detector,
betas=(0.9, 0.999),
weight_decay=0)
# place holder for GPU tensors
self.pc = torch.empty(self.opt.batch_size, 3, self.opt.input_pt_num, dtype=torch.float, device=self.opt.device)
self.intensity = torch.empty(self.opt.batch_size, 1, self.opt.input_pt_num, dtype=torch.float,
device=self.opt.device)
self.sn = torch.empty(self.opt.batch_size, 3, self.opt.input_pt_num, dtype=torch.float, device=self.opt.device)
self.node_a = torch.empty(self.opt.batch_size, 3, self.opt.node_a_num, dtype=torch.float,
device=self.opt.device)
self.node_b = torch.empty(self.opt.batch_size, 3, self.opt.node_b_num, dtype=torch.float,
device=self.opt.device)
self.P = torch.empty(self.opt.batch_size, 3, 4, dtype=torch.float, device=self.opt.device)
self.img = torch.empty(self.opt.batch_size, 3, self.opt.img_H, self.opt.img_W, dtype=torch.float,
device=self.opt.device)
self.K = torch.empty(self.opt.batch_size, 3, 3, dtype=torch.float, device=self.opt.device)
# record the train / test loss and accuracy
self.train_loss_dict = {}
self.test_loss_dict = {}
self.train_accuracy = {}
self.test_accuracy = {}
val_dataset = Dataset(opt.val_list,
phase='val',
input_shape=opt.input_shape)
trainloader = data.DataLoader(train_dataset,
batch_size=opt.train_batch_size,
shuffle=True,
num_workers=opt.num_workers)
valloader = data.DataLoader(val_dataset,
batch_size=opt.test_batch_size,
shuffle=True,
num_workers=opt.num_workers)
print('{} train iters per epoch:'.format(len(trainloader)))
############ choose loss and backbone ######################
if opt.loss == 'focal_loss':
criterion = focal_loss.FocalLoss(gamma=2)
else:
criterion = torch.nn.CrossEntropyLoss()
if opt.backbone == 'resnet18':
model = resnet.resnet_face18(use_se=opt.use_se)
elif opt.backbone == 'resnet50':
model = resnet.resnet50()
elif opt.backbone == 'resnet101':
model = resnet.resnet101()
elif opt.backbone == 'resnet152':
model = resnet.resnet152()
elif opt.backbone == 'googlenet':
model = googlenet.GoogLeNet()
if opt.metric == 'add_margin':
