Пример #1
0
                   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)

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..')
Пример #2
0
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
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)