def main():
# --------------------------------------model----------------------------------------
model = sphere20()
model = torch.nn.DataParallel(model).cuda()
print(model)
if not os.path.exists(args.save_path):
os.makedirs(args.save_path)
model.module.save(args.save_path + 'CosFace_0_checkpoint.pth')
# ------------------------------------load image---------------------------------------
train_loader = torch.utils.data.DataLoader(
ImageList(root=args.root_path, fileList=args.train_list,
transform=transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) # range [0.0, 1.0] -> [-1.0,1.0]
])),
batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True, drop_last=True)
print('length of train Dataset: ' + str(len(train_loader.dataset)))
print('Number of Classses: ' + str(args.num_class))
# --------------------------------loss function and optimizer-----------------------------
MCP = layer.MarginCosineProduct(512, args.num_class).cuda()
# MCP = layer.AngleLinear(512, args.num_class).cuda()
# MCP = torch.nn.Linear(512, args.num_class, bias=False).cuda()
criterion = torch.nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD([{'params': model.parameters()}, {'params': MCP.parameters()}],
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# ----------------------------------------train----------------------------------------
# lfw_eval.eval(args.save_path + 'CosFace_0_checkpoint.pth')
for epoch in range(1, args.epochs + 1):
# scheduler.step()
train(train_loader, model, MCP, criterion, optimizer, epoch)
model.module.save(args.save_path + 'CosFace_' + str(epoch) + '_checkpoint.pth')
lfw_eval.eval(args.save_path + 'CosFace_' + str(epoch) + '_checkpoint.pth')
print('Finished Training')
def main():
# --------------------------------------model----------------------------------------
model = net_sia.LResNet50E_IR_Sia(is_gray=args.is_gray)
model_eval = net_sia.LResNet50E_IR_Sia(is_gray=args.is_gray)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
# 512 is dimension of feature
classifier = {
'MCP': layer.MarginCosineProduct(512, args.num_class),
'AL': layer.AngleLinear(512, args.num_class),
'L': torch.nn.Linear(512, args.num_class, bias=False)
}[args.classifier_type]
classifier.load_state_dict(torch.load(args.weight_fc))
print(os.environ['CUDA_VISIBLE_DEVICES'], args.cuda)
pretrained = torch.load(args.weight_model)
pretrained_dict = pretrained['model_state_dict']
model_dict = model.state_dict()
model_eval_dict = model_eval.state_dict()
for k, v in pretrained_dict.items():
if k in model_dict:
model_dict[k].copy_(v)
del pretrained
del pretrained_dict
if args.resume:
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['model_state_dict'])
# classifier ckpt only save model info
classifier.load_state_dict(torch.load(args.resume_fc))
print(model)
model = torch.nn.DataParallel(model).to(device)
model_eval = model_eval.to(device)
classifier = classifier.to(device)
args.run_name = utils.get_run_name()
output_dir = os.path.join(args.save_path, args.run_name.split("_")[0])
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# ------------------------------------load image---------------------------------------
if args.is_gray:
train_transform = transforms.Compose([
transforms.Grayscale(),
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
transforms.Normalize(mean=(0.5, ), std=(0.5, ))
]) # gray
else:
train_transform = transforms.Compose([
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5,
0.5)) # range [0.0, 1.0] -> [-1.0,1.0]
])
valid_transform = transforms.Compose([
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5,
0.5)) # range [0.0, 1.0] -> [-1.0,1.0]
])
train_loader = torch.utils.data.DataLoader(dset.ImageList(
root=args.root_path,
fileList=args.train_list,
transform=train_transform),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=False,
drop_last=True)
val_loader = torch.utils.data.DataLoader(dset.ImageList(
root=args.root_path,
fileList=args.valid_list,
transform=valid_transform),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False,
drop_last=False)
print('length of train Database: ' + str(len(train_loader.dataset)) +
' Batches: ' + str(len(train_loader)))
print('length of valid Database: ' + str(len(val_loader.dataset)) +
' Batches: ' + str(len(val_loader)))
print('Number of Identities: ' + str(args.num_class))
# Get a batch of training data, (img, img_occ, label)
'''
inputs, inputs_occ, imgPair, targets = next(iter(train_loader))
out = torchvision.utils.make_grid(inputs)
out_occ = torchvision.utils.make_grid(inputs_occ)
mean = torch.tensor((0.5,0.5,0.5), dtype=torch.float32)
std = torch.tensor((0.5,0.5,0.5), dtype=torch.float32)
utils.imshow(out, mean, std, title=str(targets))
plt.savefig(output_dir + '/train.png')
utils.imshow(out_occ, mean, std, title=str(targets))
plt.savefig(output_dir + '/train_occ.png')
'''
#---------------------------------------params setting-----------------------------------
for name, param in model.named_parameters():
if 'layer' in name or 'conv1' in name or 'bn1' in name or 'prelu1' in name:
param.requires_grad = False
else:
param.requires_grad = True
print("Params to learn:")
params_to_update = []
params_to_stay = []
for name, param in model.named_parameters():
if param.requires_grad == True:
if 'sia' in name:
params_to_update.append(param)
print("Update \t", name)
else:
params_to_stay.append(param)
print("Stay \t", name)
for name, param in classifier.named_parameters():
param.requires_grad = True
params_to_stay.append(param)
print("Stay \t", name)
#--------------------------------loss function and optimizer-----------------------------
cfg = configurations[args.config]
criterion = torch.nn.CrossEntropyLoss().to(device)
criterion2 = torch.nn.L1Loss(reduction='mean').to(device)
optimizer = torch.optim.SGD([{
'params': params_to_stay,
'lr': 0,
'weight_decay': 0,
'momentum': 0
}, {
'params': params_to_update
}],
lr=cfg['lr'],
momentum=args.momentum,
weight_decay=args.weight_decay)
start_epoch = 1
if args.resume:
optimizer.load_state_dict(checkpoint['optim_state_dict'])
start_epoch = checkpoint['epoch']
del checkpoint
# ----------------------------------------train----------------------------------------
save_ckpt(model, 0, optimizer, output_dir +
'/CosFace_0_checkpoint.pth') # Not resumed, pretrained~
for epoch in range(start_epoch, cfg['epochs'] + 1):
train(train_loader, model, classifier, criterion, criterion2,
optimizer, epoch, cfg['step_size'], cfg['lr'])
save_ckpt(model, epoch, optimizer,
output_dir + '/CosFace_' + str(epoch) + '_checkpoint.pth')
print('Validating on valid set...')
valid(val_loader, model_eval,
output_dir + '/CosFace_' + str(epoch) + '_checkpoint.pth',
classifier, criterion, criterion2)
print('Finished Training')
def main():
# --------------------------------------model----------------------------------------
if args.network is 'sphere20':
model = net.sphere(type=20, is_gray=args.is_gray)
model_eval = net.sphere(type=20, is_gray=args.is_gray)
elif args.network is 'sphere64':
model = net.sphere(type=64, is_gray=args.is_gray)
model_eval = net.sphere(type=64, is_gray=args.is_gray)
elif args.network is 'LResNet50E_IR':
model = net.LResNet50E_IR(is_gray=args.is_gray)
model_eval = net.LResNet50E_IR(is_gray=args.is_gray)
else:
raise ValueError("NOT SUPPORT NETWORK! ")
model = torch.nn.DataParallel(model).to(device)
model_eval = model_eval.to(device)
print(model)
if not os.path.exists(args.save_path):
os.makedirs(args.save_path)
model.module.save(args.save_path + 'CosFace_0_checkpoint.pth')
# 512 is dimension of feature
classifier = {
'MCP': layer.MarginCosineProduct(512, args.num_class).to(device),
'AL': layer.AngleLinear(512, args.num_class).to(device),
'L': torch.nn.Linear(512, args.num_class, bias=False).to(device)
}[args.classifier_type]
# ------------------------------------load image---------------------------------------
if args.is_gray:
train_transform = transforms.Compose([
transforms.Grayscale(),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
transforms.Normalize(mean=(0.5, ), std=(0.5, ))
]) # gray
else:
train_transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5,
0.5)) # range [0.0, 1.0] -> [-1.0,1.0]
])
train_loader = torch.utils.data.DataLoader(ImageList(
root=args.root_path,
fileList=args.train_list,
transform=train_transform),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
drop_last=True)
print('length of train Database: ' + str(len(train_loader.dataset)))
print('Number of Identities: ' + str(args.num_class))
# --------------------------------loss function and optimizer-----------------------------
criterion = torch.nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD([{
'params': model.parameters()
}, {
'params': classifier.parameters()
}],
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# ----------------------------------------train----------------------------------------
# lfw_eval.eval(args.save_path + 'CosFace_0_checkpoint.pth')
for epoch in range(1, args.epochs + 1):
train(train_loader, model, classifier, criterion, optimizer, epoch)
model.module.save(args.save_path + 'CosFace_' + str(epoch) +
'_checkpoint.pth')
lfw_eval.eval(
model_eval,
args.save_path + 'CosFace_' + str(epoch) + '_checkpoint.pth',
args.is_gray)
print('Finished Training')
def train(img_path, label_path, model_path, net, batch_size, max_epoch,
checkpoint, lr, iscpu):
isgpu = not iscpu and torch.cuda.is_available()
if isgpu: cudnn.benchmark = True
rtime = time.time()
if net == 'FACE_NET':
model = FaceNet(Global.vecnum)
MCP = layer.MarginCosineProduct(Global.vecnum,
Global.class_out,
s=Global.mcp_size,
m=Global.mcp_dis)
if not isgpu:
model = model.cpu()
MCP = MCP.cpu()
#model = torch.nn.DataParallel(model)
torch.manual_seed(int(rtime))
else:
model = model.cuda()
#model = torch.nn.DataParallel(model).cuda()
torch.cuda.manual_seed(int(rtime))
MCP = MCP.cuda()
print(model)
kwargs = {'num_workers': 16, 'pin_memory': True} if isgpu else {}
transformList = []
transformList.append(transforms.RandomHorizontalFlip())
#transformList.append(transforms.RandomResizedCrop((Global.img_w,Global.img_h),scale=(0.07,1.0),ratio=(0.4,2.5)))
#transformList.append(transforms.ColorJitter(0.7,0.3,0.3,0.3))
transformList.append(transforms.Resize((Global.img_w, Global.img_h)))
transformList.append(
transforms.RandomCrop((Global.img_w, Global.img_h), 12))
#transformList.append(transforms.Resize((Global.img_w,Global.img_h)))
transformList.append(transforms.ToTensor())
transformList.append(transforms.Randomblack(mean=[0.0]))
transformList.append(transforms.RandomErasing(mean=[0.0]))
transformList.append(
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
transformSequence = transforms.Compose(transformList)
data = Img_Data(img_filepath=img_path,
label_filepath=label_path,
transform=transformSequence)
data_loader = DataLoader(dataset=data,
batch_size=batch_size,
shuffle=True,
**kwargs)
val_kwargs = {'num_workers': 1, 'pin_memory': True} if isgpu else {}
val_transformList = []
val_transformList.append(
transforms.Resize((Global.img_w, Global.img_h)))
val_transformList.append(transforms.ToTensor())
val_transformList.append(
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
val_transformSequence = transforms.Compose(val_transformList)
valdata = Img_Data(img_filepath=Global.val_img_path,
label_filepath=Global.val_label_path,
transform=val_transformSequence)
val_data_loader = DataLoader(dataset=valdata,
batch_size=1,
shuffle=False,
**val_kwargs)
loss = torch.nn.CrossEntropyLoss()
#optimizer = optim.SGD([{'params': model.parameters()}, {'params': MCP.parameters()}], lr=lr,momentum=0.9)
optimizer = optim.Adam([{
'params': model.parameters()
}, {
'params': MCP.parameters()
}],
lr=lr,
weight_decay=0,
amsgrad=True)
if checkpoint == True:
modelCheckpoint = torch.load(model_path)
model.load_state_dict(modelCheckpoint['state_dict'])
#optimizer.load_state_dict(modelCheckpoint['optimizer'])
MCP.load_state_dict(modelCheckpoint['MCP'])
start = modelCheckpoint['epoch'] + 1
else:
start = 0
for epoch_id in range(start, max_epoch):
oldtime = time.time()
loss_train = Train_Proc.epoch_train(model, data_loader, optimizer,
loss, MCP, isgpu)
newtime = time.time()
torch.save(
{
'epoch': epoch_id,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'MCP': MCP.state_dict()
}, './models/' + 'model-' + str(epoch_id) + '.pth.tar')
totaltime = newtime - oldtime
val_sum1, val_sum2 = Train_Proc.validate(model, val_data_loader,
isgpu)
logging.info(
'Epoch %d : training_loss = %.6f time = %d val1 = %.6f val2 = %.6f %.6f'
% (epoch_id, loss_train, totaltime, val_sum1, val_sum2,
val_sum1 - val_sum2))
with open(Global.loss_log, 'a+') as ftxt:
ftxt.write(
'\nEpoch %d : training_loss = %.6f time = %d val1 = %.6f val2 = %.6f %.6f\n'
% (epoch_id, loss_train, totaltime, val_sum1, val_sum2,
val_sum1 - val_sum2))
def main():
# ----------------------------------------load images----------------------------------------
train_loader = torch.utils.data.DataLoader(
ImageList(
root=root_path,
fileList=train_list,
transform=transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
transforms.Normalize(
mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5, 0.5)) # range [0.0, 1.0] -> [-1.0,1.0]
])),
batch_size=BatchSize,
shuffle=True,
num_workers=workers,
pin_memory=True,
drop_last=True)
print('length of train Dataset: ' + str(len(train_loader.dataset)))
print('Number of Classses: ' + str(num_class))
# ------------------------------------model--------------------------------------------
model_ft = net.sphere64a()
# # --------------load model---------------
# model_path = './checkpoints/mnface_30_checkpoints.pth'
# state_dict = torch.load(model_path)
# model_ft.load_state_dict(state_dict)
#------------------------------use gpu--------------------
if use_gpu:
model_ft = nn.DataParallel(model_ft).cuda()
# ------------------------------cosface loss and optimizer-------------------------
MCP = layer.MarginCosineProduct(512, num_class).cuda()
# MCP = layer.AngleLinear(512, args.num_class).cuda()
# MCP = torch.nn.Linear(512, args.num_class, bias=False).cuda()
criterion = torch.nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD([{
'params': model_ft.parameters()
}, {
'params': MCP.parameters()
}],
lr=lr_ori,
momentum=0.9,
weight_decay=0.0005)
for epoch in range(1, 38 + 1):
# # -------------------my loss----------------------------
# train(train_loader, model_ft, mining_loss, ce_loss, optimizer, epoch)
# model_ft.module.save(save_path + 'mnface_' + str(epoch) + '_checkpoints.pth')
# acc, pred = lfw_eval.eval(save_path + 'mnface_' + str(epoch) + '_checkpoints.pth')
#-------------------cos face--------------------------
train(train_loader, model_ft, MCP, criterion, optimizer, epoch)
model_ft.module.save(save_path + 'cosface_' + str(epoch) +
'_checkpoints.pth')
acc, pred = lfw_eval.eval(save_path + 'cosface_' + str(epoch) +
'_checkpoints.pth')
writer.add_scalar('Test/LFWAcc', acc, epoch)
print('finished training')
def main():
# --------------------------------------model----------------------------------------
# 调用 net.py 文件中的 sphere20() 网络
model = sphere20()
# DataParallel 的作用是让数据在多个 gpu 上运行
# 改为 cpu 版本（因 mac 不支持gpu运行）
# model = torch.nn.DataParallel(model).cuda()
print(model)
# 检测保存路径是否已存在，保存 checkpoint
if not os.path.exists(args.save_path):
os.makedirs(args.save_path)
model.module.save(args.save_path + 'CosFace_0_checkpoint.pth')
# ------------------------------------load image---------------------------------------
# 加载训练数据集
train_loader = torch.utils.data.DataLoader(
ImageList(
root=args.root_path,
fileList=args.train_list,
# 进行图像预处理
transform=transforms.Compose([
# 以 0.5 的概率水平翻转给定的 PIL 图像
transforms.RandomHorizontalFlip(),
# 将一个 PIL 图像（H*W*C）在 【0，255】范围内转化为 torch.Tensor(C*H*W) 在 【0.0，1.0】范围内
transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
# 使用 均值 mean 和标准差 standard deviation 来标准化数据
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)
) # range [0.0, 1.0] -> [-1.0,1.0]
])),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
drop_last=True)
# 打印 train dataset 的 length
print('length of train Dataset: ' + str(len(train_loader.dataset)))
# 打印 train dataset 的类别数目
print('Number of Classses: ' + str(args.num_class))
# --------------------------------loss function and optimizer-----------------------------
# 实现 cos face 的核心部分，但是使用了 cuda
# MCP = layer.MarginCosineProduct(512, args.num_class).cuda()
MCP = layer.MarginCosineProduct(512, args.num_class)
# MCP = layer.AngleLinear(512, args.num_class).cuda()
# MCP = torch.nn.Linear(512, args.num_class, bias=False).cuda()
# 修改为不用 cuda
# criterion = torch.nn.CrossEntropyLoss().cuda()
criterion = torch.nn.CrossEntropyLoss()
# 使用（被优化的） SGD 优化器
optimizer = torch.optim.SGD([{
'params': model.parameters()
}, {
'params': MCP.parameters()
}],
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# ----------------------------------------train----------------------------------------
# lfw_eval.eval(args.save_path + 'CosFace_0_checkpoint.pth')
# 开始训练
# 训练 epoch 次，每完整训练一次，存储一次 checkpoint
for epoch in range(1, args.epochs + 1):
# scheduler.step()
train(train_loader, model, MCP, criterion, optimizer, epoch)
model.module.save(args.save_path + 'CosFace_' + str(epoch) +
'_checkpoint.pth')
lfw_eval.eval(args.save_path + 'CosFace_' + str(epoch) +
'_checkpoint.pth')
print('Finished Training')