args = parser.parse_args()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
dataloaders, dataset_sizes, dataset = load_data(args.batch_size,
dataset=args.dataset)
model = MobileFaceNet(args.feature_dim).to(
device) # embeding size is 512 (feature vector)
print('MobileFaceNet face detection model loaded')
margin = Arcface(embedding_size=args.feature_dim,
classnum=int(dataset['train'].class_nums),
s=32.,
m=0.5).to(device)
criterion = torch.nn.CrossEntropyLoss().to(device)
optimizer_ft = optim.SGD([{
'params': model.parameters(),
'weight_decay': 5e-4
}, {
'params': margin.parameters(),
'weight_decay': 5e-4
}],
lr=0.01,
momentum=0.9,
nesterov=True)
exp_lr_scheduler = lr_scheduler.MultiStepLR(optimizer_ft,
milestones=[6, 8, 10],
gamma=0.3)
start = time.time()
## save logging and weights
train_logging_file = 'train_{}_logging.txt'.format(args.dataset)
parser.add_argument('--batch_size', type=int, default=128, help='batch size for training and evaluation')
parser.add_argument('--epoch', type=int, default=12, help='number of epoches for training')
parser.add_argument('--method', type=str, default='l2_distance',
help='methold to evaluate feature similarity, l2_distance, cos_distance')
parser.add_argument('--flip', type=str, default=True, help='if flip the image with time augmentation')
args = parser.parse_args()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
dataloaders , dataset_sizes, dataset = load_data(args.batch_size, dataset = args.dataset)
model = MobileFaceNet(args.feature_dim).to(device) # embeding size is 512 (feature vector)
print('MobileFaceNet face detection model loaded')
margin = Arcface(embedding_size=args.feature_dim, classnum=int(dataset['train'].class_nums), s=32., m=0.5).to(device)
criterion = torch.nn.CrossEntropyLoss().to(device)
optimizer_ft = optim.SGD([
{'params': model.parameters(), 'weight_decay': 5e-4},
{'params': margin.parameters(), 'weight_decay': 5e-4}], lr=0.01, momentum=0.9, nesterov=True)
exp_lr_scheduler = lr_scheduler.MultiStepLR(optimizer_ft, milestones=[6, 8, 10], gamma=0.3)
start = time.time()
## save logging and weights
train_logging_file = 'train_{}_logging.txt'.format(args.dataset)
test_logging_file = 'test_{}_logging.txt'.format(args.dataset)
save_dir = 'saving_{}_ckpt'.format(args.dataset)
if os.path.exists(save_dir):
raise NameError('model dir exists!')
os.makedirs(save_dir)
best_acc = {'LFW': 0.0, 'CFP_FP': 0.0, 'AgeDB30': 0.0}
best_iters = {'LFW': 0, 'CFP_FP': 0, 'AgeDB30': 0}
total_iters = 0