def main():
device = torch.device("cuda")
model = resnet.resnet18()
print(model)
num_classes = 13938
easy_margin = False
metric_fc = metrics.ArcMarginProduct(512,
num_classes,
s=30,
m=0.5,
easy_margin=easy_margin)
model.to(device)
model = DataParallel(model)
metric_fc.to(device)
metric_fc = DataParallel(metric_fc)
lr = 1e-1 # initial learning rate
lr_step = 10
weight_decay = 5e-4
optimizer = torch.optim.SGD([{
'params': model.parameters()
}, {
'params': metric_fc.parameters()
}],
lr=lr,
weight_decay=weight_decay)
scheduler = StepLR(optimizer, step_size=lr_step, gamma=0.1)
def main():
args = parse_args()
# Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Dataset
train_loader, test_loader, classes = cifar10_data.load_dataset(
args.dataset_dir, img_show=True)
model = resnet.resnet18()
print(model)
num_classes = 10
easy_margin = False
metric_fc = metrics.ArcMarginProduct(512,
num_classes,
s=30,
m=0.5,
easy_margin=easy_margin)
model.to(device)
model = DataParallel(model)
metric_fc.to(device)
metric_fc = DataParallel(metric_fc)
lr = 1e-1 # initial learning rate
lr_step = 10
weight_decay = 5e-4
optimizer = torch.optim.SGD([{
'params': model.parameters()
}, {
'params': metric_fc.parameters()
}],
lr=lr,
weight_decay=weight_decay)
scheduler = StepLR(optimizer, step_size=lr_step, gamma=0.1)
max_epoch = 2
for i in range(max_epoch):
scheduler.step()
model.train()
for ii, (imgs, labels) in enumerate(train_loader):
# Set batch data.
imgs, labels = imgs.to(device), labels.to(device).long()
feature = model(imgs)
output = metric_fc(feature, labels)
loss = criterion(output, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(loss)
elif opt.backbone == 'resnet101':
model = resnet.resnet101()
elif opt.backbone == 'resnet152':
model = resnet.resnet152()
elif opt.backbone == 'googlenet':
model = googlenet.GoogLeNet()
if opt.metric == 'add_margin':
metric_fc = metrics.AddMarginProduct(1024,
opt.num_classes,
s=30,
m=0.35)
elif opt.metric == 'arc_margin':
metric_fc = metrics.ArcMarginProduct(512,
opt.num_classes,
s=30,
m=0.5,
easy_margin=opt.easy_margin)
elif opt.metric == 'sphere':
metric_fc = metrics.SphereProduct(512, opt.num_classes, m=4)
else:
# metric_fc = nn.Linear(512, opt.num_classes)
metric_fc = nn.Linear(512, opt.num_classes)
############ visual_model and model_to_device ##############
tensor_board = TensorBoard(opt.train_batch_size, 3, 112, 112)
#tensor_board.visual_model(model)
model.to(device)
model = DataParallel(model)
metric_fc.to(device)
metric_fc = DataParallel(metric_fc)
# 验证集
identity_list = dataset.get_lfw_list(opt.lfw_test_list)
lfw_img_paths = [os.path.join(opt.lfw_root, each) for each in identity_list] # 所有图片的路径
# 读取训练数据集
train_dataset = dataset.Dataset(opt.train_root, opt.path_split, phase='train', input_shape=opt.input_shape)
trainloader = data.DataLoader(train_dataset,
batch_size=opt.train_batch_size,
shuffle=True,
num_workers=opt.num_workers)
opt.num_classes = len(train_dataset.classes) # 分类数量
epoch_iters = len(trainloader) # 每个epoch里iter总个数
criterion = focal_loss.FocalLoss(gamma=2)
metric_fc = metrics.ArcMarginProduct(opt.embedding, opt.num_classes, s=64, m=0.5, easy_margin=opt.easy_margin)
# 加载模型
model = mobileNetV3_MixNet.MobileNetV3_MixNet(n_class=opt.embedding, input_size=opt.input_shape[2], dropout=opt.dropout_rate)
model.to(device)
model = DataParallel(model)
metric_fc.to(device)
metric_fc = DataParallel(metric_fc)
optimizer = torch.optim.SGD([{'params': model.parameters()}, {'params': metric_fc.parameters()}],
lr=opt.lr, weight_decay=opt.weight_decay, momentum=opt.momentum)
print("epoch:{}".format(opt.max_epoch))
print("iters/epoch:{}".format(epoch_iters))
print("classes:{}".format(len(train_dataset.classes)))
print("batch_size:{}".format(opt.train_batch_size))