def fine_tune(_trian_loader, _model, n_epoch): loss_func = net_sphere.AngleLoss().cuda() optimizer = torch.optim.SGD(_model.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4) _model.cuda() _model.train() for epoch in range(1, n_epoch+1): correct = 0 train_loss = 0 total = 0 for batch_idx, (data, labels) in enumerate(_trian_loader): input_var = Variable(data, volatile=True).cuda(async=True) target_var = Variable(labels, volatile=True).cuda().long() output = _model(input_var) _, predicted = torch.max(output[0].data, 1) loss = loss_func(output, target_var) total += target_var.size(0) train_loss += loss.data[0] correct += predicted.eq(target_var.data).cpu().sum() printoneline('Te=%d Loss=%.4f | AccT=%.4f (%d/%d)' % (epoch, train_loss / (batch_idx + 1), 100.0 * correct / float(total), correct, total)) # compute gradient and do SGD step optimizer.zero_grad() loss.backward() optimizer.step() save_model(model, './%depoch.pth' % epoch) print("\n") return _model
maskNet = getattr(adversary, "MaskMan")(512) maskNet.load_state_dict(torch.load('saved_models_ce/maskNet_' + str(args.checkpoint) + '.pth')) fcNet = getattr(net_sphere, "fclayers")() # pretrainedDict = torch.load('model/sphere20a_20171020.pth') # fcDict = {k: pretrainedDict[k] for k in pretrainedDict if k in fcNet.state_dict()} fcNet.load_state_dict(torch.load('saved_models_ce/fcNet_'+ str(args.checkpoint)+ '.pth')) laplacianKernel = getKernel() # print(advNet) # net = getattr(net_sphere, "newNetwork")(net1, advNet) if torch.cuda.is_available(): featureNet.cuda() maskNet.cuda() fcNet.cuda() laplacianKernel = laplacianKernel.cuda() criterion = net_sphere.AngleLoss() optimizerFC = optim.SGD(list(featureNet.parameters()) + list(fcNet.parameters()), lr=args.lr, momentum=0.9, weight_decay=5e-4) # optimizerFeature = optim.SGD(featureNet.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4) optimizerMask = optim.SGD(maskNet.parameters(), lr = args.lr/1000, momentum=0.9, weight_decay=5e-4) criterion2 = torch.nn.CrossEntropyLoss() print('start: time={}'.format(dt())) for epoch in range(0, 50): if epoch in [0,10,15,18]: if epoch!=0: args.lr *= 0.1 optimizerFC = optim.SGD(list(featureNet.parameters()) + list(fcNet.parameters()), lr=args.lr, momentum=0.9, weight_decay=5e-4) # optimizerFeature = optim.SGD(featureNet.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4) optimizerMask = optim.SGD(maskNet.parameters(), lr = args.lr/1000, momentum=0.9, weight_decay=5e-4) if args.checkpoint >= epoch: continue # optimizerFC = optim.SGD(fcNet.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if use_gpu: print("Currently using GPU: {}".format(args.gpu)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU") print("Creating dataset: {}".format(args.dataset)) trans = transforms.Compose([transforms.Resize((200, 200)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]) train_set = dataset.UTK_DS(TRAIN_LAND_PATH,TRAIN_ROOT_PATH, args.label, trans) test_set = dataset.UTK_DS(TEST_LAND_PATH,TEST_ROOT_PATH, args.label, trans) trainloader = torch.utils.data.DataLoader(train_set, batch_size = args.bs, shuffle = True, num_workers = args.workers, drop_last = True) testloader = torch.utils.data.DataLoader(test_set, batch_size = args.bs, shuffle = True, num_workers = args.workers, drop_last = True) print("Creating model: {}".format(args.model)) if args.model == 'sphere20a': model = net_sphere.sphere20a(classnum = num_class) else: model = net_sphere.sphere64a(classnum = num_class) if use_gpu: model = nn.DataParallel(model).cuda() if os.path.isfile(model_PATH): print("\nLoading the lastest model\n") model.load_state_dict(torch.load(model_PATH)) model.eval() criterion = net_sphere.AngleLoss() lr = args.lr start_time = time.time() for epoch in range(args.max_epoch): if epoch > 0 and epoch % == 0: lr = lr * 0.1 optimizer = torch.optim.SGD(model.parameters(), lr = lr, momentum = 0.5, weight_decay = 5e-4) print("==> Epoch {}/{}".format(epoch+1, args.max_epoch)) train(model, criterion, optimizer, trainloader, use_gpu, num_class, epoch) print("==> Test") acc, err = test(model, testloader, use_gpu, num_class , epoch) print("Accuracy (%): {}\t Error rate (%): {}".format(acc, err)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) print("Finished. Total elapsed time (h:m:s): {}".format(elapsed)) torch.save(model.state_dict(), model_PATH)
train_set = mx.gluon.data.vision.ImageRecordDataset( "/home/liuhao/dataset/VGGFACE/vggface64.rec", transform=transform) train_loader = mx.gluon.data.DataLoader(dataset=train_set, batch_size=batch_size, shuffle=True, last_batch="keep") ctx = [mx.gpu(2)] my_net = net_sphere.sphere_net(classnum=8528, m=1) #mobilenet = gluon.model_zoo.vision.mobilenet1_0(pretrained=False,ctx=ctx,classes=10) my_net.initialize(init=mx.init.Xavier(), ctx=ctx) sphere_loss = net_sphere.AngleLoss() trainer = gluon.Trainer(my_net.collect_params(), 'adam', {'learning_rate': 0.5}) epoches = 30 for e in range(epoches): mean_loss = 0 my_net.get_feature = False for i, (data, label) in enumerate(train_loader): data_list = gluon.utils.split_and_load(data, ctx) label_list = gluon.utils.split_and_load(label, ctx) with autograd.record(): #training mode losses = [ sphere_loss(my_net(X), y)
def training_routine(net, n_epochs, lr, gpu, train_loader, val_loader, layer_name, embedding_size): gpu = gpu and torch.cuda.is_available() if not gpu: print('Not using GPU.') import logging logging.basicConfig(filename='train.log', level=logging.DEBUG) # criterion = nn.CrossEntropyLoss() criterion = net_sphere.AngleLoss() optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4, nesterov=True) # optimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=1e-5) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.1, patience=6) if gpu: net.cuda() # switch to train mode net.train() best_rate = 100 for i in range(n_epochs): tic = time.time() train_prediction = [] train_observed = [] train_loss_avg = 0 train_loss_epochs = 0 for j, (train_labels, train_data) in enumerate(train_loader): if gpu: train_labels, train_data = train_labels.cuda( ), train_data.cuda() # forward pass train_output = net(train_data) train_loss = criterion(train_output, train_labels) # backward pass and optimization optimizer.zero_grad() train_loss.backward() optimizer.step() train_loss_avg += train_loss.cpu().detach().numpy() train_loss_epochs += train_loss.cpu().detach().numpy() # train_output = train_output.cpu().argmax(dim=1).detach().numpy() train_prediction.append(train_output) train_labels = np.array(train_labels.cpu().numpy()) train_observed.append(train_labels) torch.cuda.empty_cache() # training print batch_print = 40 if j % batch_print == 0 and j != 0: t = 'At {:.0f}% of epoch {}'.format( j * train_loader.batch_size / train_loader.dataset.num_entries * 100, i) print(t) logging.info(t) # train_accuracy = np.array(train_output == train_labels).mean() t = "Training loss : {}".format(train_loss_epochs / batch_print) train_loss_epochs = 0 print(t) logging.info(t) # t = "Training accuracy {}:".format(train_accuracy) # print(t) # logging.info(t) t = '--------------------------------------------' print(t) logging.info(t) scheduler.step(train_loss_avg / j) # every 1 epochs, print validation statistics epochs_print = 10 if i % epochs_print == 0 and not i == 0: with torch.no_grad(): t = "######### Epoch {} #########".format(i) print(t) logging.info(t) # compute the accuracy of the prediction # train_prediction = np.concatenate(train_prediction) # train_observed = np.concatenate(train_observed) # train_accuracy = (train_prediction == train_observed).mean() # Now for the validation set val_prediction = [] val_observed = [] enrol = {} test = {} for j, (trial, val_labels, val_enrol, val_test) in (enumerate(val_loader)): if gpu: val_labels, val_enrol, val_test = val_labels.cuda( ), val_enrol.cuda(), val_test.cuda() key_enrol_array, key_test_array = trial[:, 0], trial[:, 1] embedding_test = [] embedding_enrol = [] for t in range(len(key_enrol_array)): key_enrol = key_enrol_array[t] key_test = key_test_array[t] if key_test not in test: test[key_test] = extract_embedding( val_test[t].unsqueeze(0), net, layer_name, (len(val_test[t]), embedding_size)) embedding_test.append(test[key_test]) if key_enrol not in enrol: enrol[key_enrol] = extract_embedding( val_enrol[t].unsqueeze(0), net, layer_name, (len(val_enrol[t]), embedding_size)) embedding_enrol.append(enrol[key_enrol]) embedding_enrol = torch.cat(embedding_enrol) embedding_test = torch.cat(embedding_test) cos = torch.nn.CosineSimilarity() val_output = cos(embedding_test, embedding_enrol) val_prediction.append(val_output) val_labels = val_labels.cpu().numpy() val_observed.append(val_labels) val_prediction = np.concatenate(val_prediction) val_observed = np.concatenate(val_observed) # compute the accuracy of the prediction val_eed = utils.EER(val_observed, val_prediction) # t = "Training accuracy : {}".format(train_accuracy) # print(t) # logging.info(t) t = "Validation EER {}:".format(val_eed) print(t) logging.info(t) toc = time.time() t = "Took: {}".format((toc - tic) / epochs_print) print(t) logging.info(t) t = '--------------------------------------------' print(t) logging.info(t) if best_rate > val_eed[0]: save_model(net, 'model.torch') best_rate = val_eed[0] net = net.cpu() return net