def main(args):
dist.init_process_group(backend='nccl', init_method='env://')
local_rank = args.local_rank
torch.cuda.set_device(local_rank)
rank = dist.get_rank()
world_size = dist.get_world_size()
if not os.path.exists(cfg.output) and rank is 0:
os.makedirs(cfg.output)
else:
time.sleep(2)
log_root = logging.getLogger()
init_logging(log_root, rank, cfg.output)
trainset = MXFaceDataset(root_dir=cfg.rec, local_rank=local_rank)
train_sampler = torch.utils.data.distributed.DistributedSampler(
trainset, shuffle=True)
train_loader = DataLoaderX(local_rank=local_rank,
dataset=trainset,
batch_size=cfg.batch_size,
sampler=train_sampler,
num_workers=0,
pin_memory=True,
drop_last=True)
dropout = 0.4 if cfg.dataset is "webface" else 0
backbone = eval("backbones.{}".format(args.network))(
False, dropout=dropout, fp16=cfg.fp16).to(local_rank)
if args.resume:
try:
backbone_pth = os.path.join(cfg.output, "backbone.pth")
backbone.load_state_dict(
torch.load(backbone_pth,
map_location=torch.device(local_rank)))
if rank is 0:
logging.info("backbone resume successfully!")
except (FileNotFoundError, KeyError, IndexError, RuntimeError):
logging.info("resume fail, backbone init successfully!")
for ps in backbone.parameters():
dist.broadcast(ps, 0)
backbone = torch.nn.parallel.DistributedDataParallel(
module=backbone, broadcast_buffers=False, device_ids=[local_rank])
backbone.train()
margin_softmax = eval("losses.{}".format(args.loss))()
module_partial_fc = PartialFC(rank=rank,
local_rank=local_rank,
world_size=world_size,
resume=args.resume,
batch_size=cfg.batch_size,
margin_softmax=margin_softmax,
num_classes=cfg.num_classes,
sample_rate=cfg.sample_rate,
embedding_size=cfg.embedding_size,
prefix=cfg.output)
opt_backbone = torch.optim.SGD(params=[{
'params': backbone.parameters()
}],
lr=cfg.lr / 512 * cfg.batch_size *
world_size,
momentum=0.9,
weight_decay=cfg.weight_decay)
opt_pfc = torch.optim.SGD(params=[{
'params': module_partial_fc.parameters()
}],
lr=cfg.lr / 512 * cfg.batch_size * world_size,
momentum=0.9,
weight_decay=cfg.weight_decay)
scheduler_backbone = torch.optim.lr_scheduler.LambdaLR(
optimizer=opt_backbone, lr_lambda=cfg.lr_func)
scheduler_pfc = torch.optim.lr_scheduler.LambdaLR(optimizer=opt_pfc,
lr_lambda=cfg.lr_func)
start_epoch = 0
total_step = int(
len(trainset) / cfg.batch_size / world_size * cfg.num_epoch)
if rank is 0: logging.info("Total Step is: %d" % total_step)
callback_verification = CallBackVerification(2000, rank, cfg.val_targets,
cfg.rec)
callback_logging = CallBackLogging(50, rank, total_step, cfg.batch_size,
world_size, None)
callback_checkpoint = CallBackModelCheckpoint(rank, cfg.output)
loss = AverageMeter()
global_step = 0
grad_scaler = MaxClipGradScaler(
cfg.batch_size, 128 *
cfg.batch_size, growth_interval=100) if cfg.fp16 else None
for epoch in range(start_epoch, cfg.num_epoch):
train_sampler.set_epoch(epoch)
for step, (img, label) in enumerate(train_loader):
global_step += 1
features = F.normalize(backbone(img))
x_grad, loss_v = module_partial_fc.forward_backward(
label, features, opt_pfc)
if cfg.fp16:
features.backward(grad_scaler.scale(x_grad))
grad_scaler.unscale_(opt_backbone)
clip_grad_norm_(backbone.parameters(), max_norm=5, norm_type=2)
grad_scaler.step(opt_backbone)
grad_scaler.update()
else:
features.backward(x_grad)
clip_grad_norm_(backbone.parameters(), max_norm=5, norm_type=2)
opt_backbone.step()
opt_pfc.step()
module_partial_fc.update()
opt_backbone.zero_grad()
opt_pfc.zero_grad()
loss.update(loss_v, 1)
callback_logging(global_step, loss, epoch, cfg.fp16, grad_scaler)
callback_verification(global_step, backbone)
callback_checkpoint(global_step, backbone, module_partial_fc)
scheduler_backbone.step()
scheduler_pfc.step()
dist.destroy_process_group()
def main(args):
world_size = int(1.0)
rank = int(0.0)
local_rank = args.local_rank
if not os.path.exists(cfg.output):
os.makedirs(cfg.output)
else:
time.sleep(2)
if not os.path.exists(cfg.output):
os.makedirs(cfg.output)
else:
time.sleep(2)
writer = LogWriter(logdir=cfg.logdir)
trainset = MXFaceDataset(root_dir=cfg.rec)
train_loader = DataLoader(dataset=trainset,
batch_size=cfg.batch_size,
shuffle=True,
drop_last=True,
num_workers=0)
dropout = 0.4 if cfg.dataset == "webface" else 0
backbone = eval("backbones.{}".format(args.network))(False,
dropout=0.5,
fp16=False)
backbone.train()
clip_by_norm = ClipGradByNorm(5.0)
margin_softmax = eval("losses.{}".format(args.loss))()
module_partial_fc = PartialFC(rank=0,
local_rank=0,
world_size=1,
resume=0,
batch_size=cfg.batch_size,
margin_softmax=margin_softmax,
num_classes=cfg.num_classes,
sample_rate=cfg.sample_rate,
embedding_size=cfg.embedding_size,
prefix=cfg.output)
scheduler_backbone = paddle.optimizer.lr.LambdaDecay(learning_rate=cfg.lr /
512 * cfg.batch_size,
lr_lambda=cfg.lr_func,
verbose=True)
opt_backbone = paddle.optimizer.SGD(parameters=backbone.parameters(),
learning_rate=scheduler_backbone,
weight_decay=cfg.weight_decay,
grad_clip=clip_by_norm)
scheduler_pfc = paddle.optimizer.lr.LambdaDecay(learning_rate=cfg.lr /
512 * cfg.batch_size,
lr_lambda=cfg.lr_func,
verbose=True)
opt_pfc = paddle.optimizer.SGD(parameters=module_partial_fc.parameters(),
learning_rate=scheduler_pfc,
weight_decay=cfg.weight_decay,
grad_clip=clip_by_norm)
start_epoch = 0
total_step = int(
len(trainset) / cfg.batch_size / world_size * cfg.num_epoch)
if rank == 0:
print("Total Step is: %d" % total_step)
callback_verification = CallBackVerification(2000, rank, cfg.val_targets,
cfg.rec)
callback_logging = CallBackLogging(100, rank, total_step, cfg.batch_size,
world_size, writer)
callback_checkpoint = CallBackModelCheckpoint(rank, cfg.output)
loss = AverageMeter()
global_step = 0
grad_scaler = MaxClipGradScaler(
cfg.batch_size, 128 *
cfg.batch_size, growth_interval=100) if cfg.fp16 else None
for epoch in range(start_epoch, cfg.num_epoch):
for step, (img, label) in enumerate(train_loader):
label = label.flatten()
global_step += 1
features = F.normalize(backbone(img))
x_grad, loss_v = module_partial_fc.forward_backward(
label, features, opt_pfc)
if cfg.fp16:
scaled = grad_scaler.scale(x_grad)
(features.multiply(scaled)).backward()
grad_scaler._unscale(opt_backbone)
grad_scaler.minimize(opt_backbone, scaled)
else:
(features.multiply(x_grad)).backward()
opt_backbone.step()
opt_pfc.step()
module_partial_fc.update()
opt_backbone.clear_gradients()
opt_pfc.clear_gradients()
loss.update(loss_v, 1)
callback_logging(global_step, loss, epoch, cfg.fp16, grad_scaler)
callback_verification(global_step, backbone)
callback_checkpoint(global_step, backbone, module_partial_fc)
scheduler_backbone.step()
scheduler_pfc.step()
writer.close()
def main(args):
world_size = int(1.0)
rank = int(0.0)
if not os.path.exists(args.output):
os.makedirs(args.output)
else:
time.sleep(2)
writer = LogWriter(logdir=args.logdir)
trainset = CommonDataset(root_dir=cfg.data_dir, label_file=cfg.file_list, is_bin=args.is_bin)
train_loader = DataLoader(
dataset=trainset,
batch_size=args.batch_size,
shuffle=True,
drop_last=True,
num_workers=0)
backbone = eval("backbones.{}".format(args.network))()
backbone.train()
clip_by_norm = ClipGradByNorm(5.0)
margin_softmax = eval("losses.{}".format(args.loss))()
module_partial_fc = PartialFC(
rank=0,
world_size=1,
resume=0,
batch_size=args.batch_size,
margin_softmax=margin_softmax,
num_classes=cfg.num_classes,
sample_rate=cfg.sample_rate,
embedding_size=args.embedding_size,
prefix=args.output)
scheduler_backbone_decay = paddle.optimizer.lr.LambdaDecay(
learning_rate=args.lr, lr_lambda=cfg.lr_func, verbose=True)
scheduler_backbone = paddle.optimizer.lr.LinearWarmup(
learning_rate=scheduler_backbone_decay,
warmup_steps=cfg.warmup_epoch,
start_lr=0,
end_lr=args.lr / 512 * args.batch_size,
verbose=True)
opt_backbone = paddle.optimizer.Momentum(
parameters=backbone.parameters(),
learning_rate=scheduler_backbone,
momentum=0.9,
weight_decay=args.weight_decay,
grad_clip=clip_by_norm)
scheduler_pfc_decay = paddle.optimizer.lr.LambdaDecay(
learning_rate=args.lr, lr_lambda=cfg.lr_func, verbose=True)
scheduler_pfc = paddle.optimizer.lr.LinearWarmup(
learning_rate=scheduler_pfc_decay,
warmup_steps=cfg.warmup_epoch,
start_lr=0,
end_lr=args.lr / 512 * args.batch_size,
verbose=True)
opt_pfc = paddle.optimizer.Momentum(
parameters=module_partial_fc.parameters(),
learning_rate=scheduler_pfc,
momentum=0.9,
weight_decay=args.weight_decay,
grad_clip=clip_by_norm)
start_epoch = 0
total_step = int(
len(trainset) / args.batch_size / world_size * cfg.num_epoch)
if rank == 0:
print("Total Step is: %d" % total_step)
callback_verification = CallBackVerification(2000, rank, cfg.val_targets,
cfg.data_dir)
callback_logging = CallBackLogging(10, rank, total_step, args.batch_size,
world_size, writer)
callback_checkpoint = CallBackModelCheckpoint(rank, args.output,
args.network)
loss = AverageMeter()
global_step = 0
for epoch in range(start_epoch, cfg.num_epoch):
for step, (img, label) in enumerate(train_loader):
label = label.flatten()
global_step += 1
sys.stdout.flush()
features = F.normalize(backbone(img))
x_grad, loss_v = module_partial_fc.forward_backward(
label, features, opt_pfc)
sys.stdout.flush()
(features.multiply(x_grad)).backward()
sys.stdout.flush()
opt_backbone.step()
opt_pfc.step()
module_partial_fc.update()
opt_backbone.clear_gradients()
opt_pfc.clear_gradients()
sys.stdout.flush()
lr_backbone_value = opt_backbone._global_learning_rate().numpy()[0]
lr_pfc_value = opt_backbone._global_learning_rate().numpy()[0]
loss.update(loss_v, 1)
callback_logging(global_step, loss, epoch, lr_backbone_value,
lr_pfc_value)
sys.stdout.flush()
callback_verification(global_step, backbone)
callback_checkpoint(global_step, backbone, module_partial_fc)
scheduler_backbone.step()
scheduler_pfc.step()
writer.close()
def main(args):
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
cfg = get_config(args.config)
try:
world_size = int(os.environ['WORLD_SIZE'])
rank = int(os.environ['RANK'])
dist.init_process_group('nccl')
except KeyError:
world_size = 1
rank = 0
dist.init_process_group(backend='nccl',
init_method="tcp://127.0.0.1:12584",
rank=rank,
world_size=world_size)
local_rank = args.local_rank
torch.cuda.set_device(local_rank)
os.makedirs(cfg.output, exist_ok=True)
init_logging(rank, cfg.output)
if cfg.rec == "synthetic":
train_set = SyntheticDataset(local_rank=local_rank)
else:
train_set = MXFaceDataset(root_dir=cfg.rec, local_rank=local_rank)
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_set, shuffle=True)
train_loader = DataLoaderX(local_rank=local_rank,
dataset=train_set,
batch_size=cfg.batch_size,
sampler=train_sampler,
num_workers=2,
pin_memory=True,
drop_last=True)
backbone = get_model(cfg.network,
dropout=0.0,
fp16=cfg.fp16,
num_features=cfg.embedding_size).to(local_rank)
summary(backbone, input_size=(3, 112, 112))
exit()
if cfg.resume:
try:
backbone_pth = os.path.join(cfg.output, "backbone.pth")
backbone.load_state_dict(
torch.load(backbone_pth,
map_location=torch.device(local_rank)))
if rank == 0:
logging.info("backbone resume successfully!")
except (FileNotFoundError, KeyError, IndexError, RuntimeError):
if rank == 0:
logging.info("resume fail, backbone init successfully!")
backbone = torch.nn.parallel.DistributedDataParallel(
module=backbone, broadcast_buffers=False, device_ids=[local_rank])
backbone.train()
if cfg.loss == 'magface':
margin_softmax = losses.get_loss(cfg.loss, lambda_g=cfg.lambda_g)
elif cfg.loss == 'mag_cosface':
margin_softmax = losses.get_loss(cfg.loss)
else:
margin_softmax = losses.get_loss(cfg.loss,
s=cfg.s,
m1=cfg.m1,
m2=cfg.m2,
m3=cfg.m3)
module_partial_fc = PartialFC(rank=rank,
local_rank=local_rank,
world_size=world_size,
resume=cfg.resume,
batch_size=cfg.batch_size,
margin_softmax=margin_softmax,
num_classes=cfg.num_classes,
sample_rate=cfg.sample_rate,
embedding_size=cfg.embedding_size,
prefix=cfg.output)
opt_backbone = torch.optim.SGD(params=[{
'params': backbone.parameters()
}],
lr=cfg.lr / 512 * cfg.batch_size *
world_size,
momentum=0.9,
weight_decay=cfg.weight_decay)
opt_pfc = torch.optim.SGD(params=[{
'params': module_partial_fc.parameters()
}],
lr=cfg.lr / 512 * cfg.batch_size * world_size,
momentum=0.9,
weight_decay=cfg.weight_decay)
num_image = len(train_set)
total_batch_size = cfg.batch_size * world_size
cfg.warmup_step = num_image // total_batch_size * cfg.warmup_epoch
cfg.total_step = num_image // total_batch_size * cfg.num_epoch
def lr_step_func(current_step):
cfg.decay_step = [
x * num_image // total_batch_size for x in cfg.decay_epoch
]
if current_step < cfg.warmup_step:
return current_step / cfg.warmup_step
else:
return 0.1**len([m for m in cfg.decay_step if m <= current_step])
scheduler_backbone = torch.optim.lr_scheduler.LambdaLR(
optimizer=opt_backbone, lr_lambda=lr_step_func)
scheduler_pfc = torch.optim.lr_scheduler.LambdaLR(optimizer=opt_pfc,
lr_lambda=lr_step_func)
for key, value in cfg.items():
num_space = 25 - len(key)
logging.info(": " + key + " " * num_space + str(value))
val_target = cfg.val_targets
callback_verification = CallBackVerification(2000, rank, val_target,
cfg.rec)
callback_logging = CallBackLogging(50, rank, cfg.total_step,
cfg.batch_size, world_size, None)
callback_checkpoint = CallBackModelCheckpoint(rank, cfg.output)
loss = AverageMeter()
start_epoch = 0
global_step = 0
grad_amp = MaxClipGradScaler(
cfg.batch_size, 128 *
cfg.batch_size, growth_interval=100) if cfg.fp16 else None
for epoch in range(start_epoch, cfg.num_epoch):
train_sampler.set_epoch(epoch)
for step, (img, label) in enumerate(train_loader):
global_step += 1
x = backbone(img)
features = F.normalize(x)
x_grad, loss_v = module_partial_fc.forward_backward(
label, features, opt_pfc, x)
if cfg.fp16:
features.backward(grad_amp.scale(x_grad))
grad_amp.unscale_(opt_backbone)
clip_grad_norm_(backbone.parameters(), max_norm=5, norm_type=2)
grad_amp.step(opt_backbone)
grad_amp.update()
else:
features.backward(x_grad)
clip_grad_norm_(backbone.parameters(), max_norm=5, norm_type=2)
opt_backbone.step()
opt_pfc.step()
module_partial_fc.update()
opt_backbone.zero_grad()
opt_pfc.zero_grad()
loss.update(loss_v, 1)
callback_logging(global_step, loss, epoch, cfg.fp16,
scheduler_backbone.get_last_lr()[0], grad_amp)
callback_verification(global_step, backbone)
scheduler_backbone.step()
scheduler_pfc.step()
callback_checkpoint(global_step, backbone, module_partial_fc)
callback_verification('last', backbone)
dist.destroy_process_group()