def net_config_classify(image, label, model, args):
out = createmodel(image, model, args)
if args.loss_name == "softmax":
metricloss = SoftmaxLoss(class_dim=args.class_dim, )
elif args.loss_name == "arcmargin":
metricloss = ArcMarginLoss(
class_dim=args.class_dim,
margin=args.arc_margin,
scale=args.arc_scale,
easy_margin=args.arc_easy_margin,
)
cost, logit = metricloss.loss(out, label)
avg_cost = fluid.layers.mean(x=cost)
acc_top1 = fluid.layers.accuracy(input=logit, label=label, k=1)
acc_top5 = fluid.layers.accuracy(input=logit, label=label, k=5)
return [avg_cost, acc_top1, acc_top5]
def net_config(image, label, model, args):
assert args.model in model_list, "{} is not in lists: {}".format(args.model,
model_list)
out = model.net(input=image, embedding_size=args.embedding_size)
if args.loss_name == "softmax":
metricloss = SoftmaxLoss(class_dim=args.class_dim, )
elif args.loss_name == "arcmargin":
metricloss = ArcMarginLoss(
class_dim=args.class_dim,
margin=args.arc_margin,
scale=args.arc_scale,
easy_margin=args.arc_easy_margin, )
cost, logit = metricloss.loss(out, label)
avg_cost = fluid.layers.mean(x=cost)
acc_top1 = fluid.layers.accuracy(input=logit, label=label, k=1)
acc_top5 = fluid.layers.accuracy(input=logit, label=label, k=5)
return avg_cost, acc_top1, acc_top5, out
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=256)
parser.add_argument('--loss', default='softmax_loss',
choices=['softmax_loss', 'center_loss', 'sphere_face_loss', 'cos_face_loss', 'arc_face_loss'])
parser.add_argument('--viz', default='vizs')
parser.add_argument('--epochs', type=int, default=30)
parser.add_argument('--lr', type=float, default=0.001)
args = parser.parse_args()
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 8, 'pin_memory': True} if use_cuda else {}
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_loader = DataLoader(
datasets.MNIST('../data', train=True, download=True, transform=transform),
batch_size=args.batch_size, shuffle=True, **kwargs)
test_loader = DataLoader(
datasets.MNIST('../data', train=False, transform=transform),
batch_size=512, shuffle=True, **kwargs)
model = Net().to(device)
if args.loss == 'center_loss':
criterion = CenterLoss().to(device)
center_optimizer = optim.SGD([criterion.centers], lr=args.lr, momentum=0.9)
elif args.loss == 'sphere_face_loss':
criterion = SphereFaceLoss().to(device)
elif args.loss == 'cos_face_loss':
criterion = CosFaceLoss(s=7, m=0.2).to(device)
elif args.loss == 'softmax_loss':
criterion = SoftmaxLoss().to(device)
elif args.loss == 'arc_face_loss':
criterion = ArcFaceLoss().to(device)
optimizer = optim.SGD([{'params': model.parameters()}, {'params': criterion.fc.parameters()}],
lr=args.lr, momentum=0.9)
for epoch in range(1, args.epochs + 1):
model.train()
embeddings = []
labels = []
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
if args.loss == 'center_loss':
center_optimizer.zero_grad()
embedding = model(data)
loss = criterion(embedding, target)
loss.backward()
optimizer.step()
if args.loss == 'center_loss':
center_optimizer.step()
embeddings.append(embedding)
labels.append(target)
if batch_idx % 100 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
embeddings = torch.cat(embeddings, 0).cpu().detach().numpy()
labels = torch.cat(labels, 0).cpu().detach().numpy()
acc = val(model, criterion, device, test_loader)
visualize(args.viz, args.loss, embeddings, labels, epoch, acc)
print('Creating gif...')
create_gif('./%s/gifs/%s.gif' % (args.viz, args.loss),
'./%s/%s' % (args.viz, args.loss), 0.2)
print('Done')