def main_stage1():
print(f"\nStart training ...\n")
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
print('==> Building model..')
net = BuildNet(backbone=args.arch, num_classes=args.train_class_num, embed_dim=args.embed_dim)
net = net.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
criterion = PSoftmaxLoss()
criterion = criterion.to(device)
optimizer = torch.optim.SGD(net.parameters(), lr=args.stage1_lr, momentum=0.9, weight_decay=5e-4)
if args.stage1_resume:
# Load checkpoint.
if os.path.isfile(args.stage1_resume):
print('==> Resuming from checkpoint..')
checkpoint = torch.load(args.stage1_resume)
net.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch']
logger = Logger(os.path.join(args.checkpoint, 'log_stage1.txt'), resume=True)
else:
print("=> no checkpoint found at '{}'".format(args.stage1_resume))
else:
logger = Logger(os.path.join(args.checkpoint, 'log_stage1.txt'))
logger.set_names(['Epoch', 'Train Loss', 'Train Acc.', "Test F1", 'threshold'])
if not args.evaluate:
for epoch in range(start_epoch, args.stage1_es):
adjust_learning_rate(optimizer, epoch, args.stage1_lr,
factor=args.stage1_lr_factor, step=args.stage1_lr_step)
print('\nEpoch: %d | Learning rate: %f ' % (epoch + 1, optimizer.param_groups[0]['lr']))
train_out = train(net, trainloader, optimizer, criterion, device)
save_model(net, optimizer, epoch, os.path.join(args.checkpoint, 'stage_1_last_model.pth'))
test_out = test(net, testloader, device)
logger.append([epoch + 1, train_out["train_loss"], train_out["accuracy"],
test_out["best_F1"], test_out["best_thres"]])
logger.close()
print(f"\nFinish training...\n")
print("===> Evaluating ...")
test_with_hist(net, testloader, device, name="stage1_test")
mid_dict = stage_valmixup(net, trainloader, device)
return {
"net": net,
"mid_known": mid_dict["mid_known"],
"mid_unknown": mid_dict["mid_unknown"]
}
batch_size=args.bs,
shuffle=True,
num_workers=4)
loss_Dict = {
"CenterLoss":
CenterLoss(centerloss_weight=args.centerloss_weight,
num_classes=args.train_class_num),
"SoftmaxLoss":
SoftmaxLoss(),
"ArcFaceLoss":
ArcFaceLoss(scaling=args.scaling, m=args.m),
"NormFaceLoss":
NormFaceLoss(scaling=args.scaling),
"PSoftmaxLoss":
PSoftmaxLoss()
}
criterion = loss_Dict[args.loss]
criterion = criterion.to(device)
def main():
print(f"\nStart training ...\n")
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
print('==> Building model..')
net = BuildNet(backbone=args.arch,
num_classes=args.train_class_num,
embed_dim=args.embed_dim)
net = net.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
])
trainset = MNIST(root='../../data', train=True, download=True, transform=transform,
train_class_num=args.train_class_num, test_class_num=args.test_class_num,
includes_all_train_class=args.includes_all_train_class)
testset = MNIST(root='../../data', train=False, download=True, transform=transform,
train_class_num=args.train_class_num, test_class_num=args.test_class_num,
includes_all_train_class=args.includes_all_train_class)
# data loader
trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.bs, shuffle=True, num_workers=4)
testloader = torch.utils.data.DataLoader(testset, batch_size=args.bs, shuffle=False, num_workers=4)
loss_Dict = {"CenterLoss": CenterLoss(centerloss_weight=args.centerloss_weight, num_classes=args.train_class_num),
"SoftmaxLoss": SoftmaxLoss(),
"ArcFaceLoss": ArcFaceLoss(scaling=args.scaling, m=args.m),
"NormFaceLoss": NormFaceLoss(scaling=args.scaling),
"PSoftmaxLoss": PSoftmaxLoss()}
criterion = loss_Dict[args.loss]
criterion = criterion.to(device)
def main():
print(f"\nStart training ...\n")
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
print('==> Building model..')
net = BuildNet(backbone=args.arch, num_classes=args.train_class_num, embed_dim=args.embed_dim)
net = net.to(device)
if device == 'cuda':
net = torch.nn.DataParallel(net)
cudnn.benchmark = True
optimizer = torch.optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)