现在我们开始看一下配置文件是怎样的,打开config.py可以看到,参数有很多,我们主要讲一下我们所调的,没有注释的默认就好。
parent_parser = argparse.ArgumentParser(add_help=False)
#数据集路径
parent_parser.add_argument('--data_root', type=str, default='')
#训练数据集的路径
parent_parser.add_argument('--train_data_path', type=str, default='faces_emore/imgs')
#验证集的路径
parent_parser.add_argument('--val_data_path', type=str, default='faces_emore')
#训练数据集
parent_parser.add_argument('--train_data_subset', action='store_true')
parent_parser.add_argument('--prefix', type=str, default='default')
#使用多少个gpu进行训练
parent_parser.add_argument('--gpus', type=int, default=1, help='how many gpus')
#用多卡多机进行训练
parent_parser.add_argument('--distributed_backend', type=str, default='ddp', choices=('dp', 'ddp', 'ddp2'),)
parent_parser.add_argument('--use_16bit', action='store_true', help='if true uses 16 bit precision')
#训练的一个总epoch
parent_parser.add_argument('--epochs', default=26, type=int, metavar='N', help='number of total epochs to run')
parent_parser.add_argument('--seed', type=int, default=42, help='seed for initializing training.')
#训练一个epoch所用到的图片,这里面是指一共训练的多少张图片,如果用多卡的话,就要一张卡可以训练的图片数量乘以几张卡。
#例如两张卡的话,一张卡能训256张图,两张卡就是2*256=512,那么下面就填512。
parent_parser.add_argument('--batch_size', default=256, type=int,
help='mini-batch size (default: 256), this is the total '
'batch size of all GPUs on the current node when '
'using Data Parallel or Distributed Data Parallel')
#学习率
parent_parser.add_argument('--lr',help='learning rate',default=0.002, type=float)
#在epoch=12、20、24的时候进行一个学习率的下降
parent_parser.add_argument('--lr_milestones', default='12,20,24', type=str, help='epochs for reducing LR')
parent_parser.add_argument('--lr_gamma', default=0.1, type=float, help='multiply when reducing LR')
#线程数量
parent_parser.add_argument('--num_workers', default=36, type=int)
parent_parser.add_argument('--fast_dev_run', dest='fast_dev_run', action='store_true')
parent_parser.add_argument('--evaluate', action='store_true', help='use with start_from_model_statedict')
parent_parser.add_argument('--resume_from_checkpoint', type=str, default='')
parent_parser.add_argument('--start_from_model_statedict', type=str, default='')
#使用ir_18作为主干网络
parser.add_argument('--arch', default='ir_18')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M')
parser.add_argument('--weight_decay', default=1e-4, type=float)
parser.add_argument('--head', default='adaface', type=str, choices=('adaface'))
parser.add_argument('--m', default=0.4, type=float)
parser.add_argument('--h', default=0.333, type=float)
parser.add_argument('--s', type=float, default=64.0)
parser.add_argument('--t_alpha', default=0.01, type=float)
parser.add_argument('--low_res_augmentation_prob', default=0.2, type=float)
parser.add_argument('--crop_augmentation_prob', default=0.2, type=float)
parser.add_argument('--photometric_augmentation_prob', default=0.2, type=float)
parser.add_argument('--accumulate_grad_batches', type=int, default=1)
parser.add_argument('--test_run', action='store_true')
parser.add_argument('--save_all_models', action='store_true')
1、参数设定好之后,我们就可以开始模型的训练了,运行一下main.py函数,即可开始训练。
2、下图为训练成功时的样子。由于我只是师范一下,我只用了一张卡来训,训练成功后我就kill掉了。
3、最后模型保存的位置在目录下的,experiments文件中,只保存最好的一个模型,如果需要保存全部模型,自行修改一下代码即可。
好了,在这里我们模型已经训练出来了,接下来我们看一下测试效果如何。
首先我们看一下推理的代码inference.py。这里我修改了一下源码,直接检测用前面所说的libface算法检测到人脸之后,保存改头像图片,之后进行一个预测,因此只需要改一下图片的输入就可以直接预测了,不需要再经过一个人脸检测的模块,这样速度也会上来。这里我把我自己的代码给分享出来,大家可以参考一下。
import net
import torch
import os
import cv2
from face_alignment import align
import numpy as np
adaface_models = {
'ir_18':r".\experiments\default_08-01_1\epoch=24-step=142174.ckpt".replace('\\','/'),
}
def load_pretrained_model(architecture='ir_18'):
# load model and pretrained statedict
assert architecture in adaface_models.keys()
model = net.build_model(architecture)
statedict = torch.load(adaface_models[architecture])['state_dict']
model_statedict = {key[6:]:val for key, val in statedict.items() if key.startswith('model.')}
model.load_state_dict(model_statedict)
model.eval()
return model
def to_input(pil_rgb_image):
np_img = np.array(pil_rgb_image)
brg_img = ((np_img[:,:,::-1] / 255.) - 0.5) / 0.5
tensor = torch.tensor([brg_img.transpose(2,0,1)]).float()
return tensor
if __name__ == '__main__':
model = load_pretrained_model('ir_18')
feature, norm = model(torch.randn(1,3,112,112))
test_image_path = 'face_alignment/test'
features = []
for fname in sorted(os.listdir(test_image_path)):
# img = cv2.imread(os.path.join(test_image_path, fname))
# img = torch.from_numpy(img)
img = cv2.imread(os.path.join(test_image_path, fname))
img = np.expand_dims(img, axis=0)
img = img.astype(np.float32)
img = img.transpose(0, 3, 1, 2)
img = torch.from_numpy(img)
# aligned_rgb_img = align.get_aligned_face(path)
# bgr_tensor_input = to_input(aligned_rgb_img)
feature, _ = model(img)
features.append(feature)
similarity_scores = torch.cat(features) @ torch.cat(features).T
print(similarity_scores)
这里我给大家讲述一下这个输出结果,我们看做这个矩阵为3*3,行就是1 2 3,竖也是1 2 3,这样我们可以看做,第一张图片跟第一张图片的相似度为 1 ,因为是自己跟自己对比,肯定是1,然后第一张图片跟第二张图片对比的相似度为0.4200,第一张图片跟第三张图片对比的相似度为0.2764,如此类推……0.3713的相似度为第二张图片跟第三张图片对比的结果。