def __init__(self, conf, inference=False):
self.conf = conf
self.logger = self.get_logger()
if self.conf.model.format == 'res50':
self.model = resnet_face50(use_se=True)
elif self.conf.model.format == 'res101':
self.model = resnet_face101(use_se=True)
else:
raise ValueError
self.model = torch.nn.DataParallel(self.model).cuda()
if not inference:
self.milestones = conf.train.milestones
self.loader = get_train_loader(conf)
self.step = 0
self.optimizer = optim.SGD(list(self.model.parameters()),
lr=conf.train.lr,
momentum=conf.train.momentum,
weight_decay=0.0005)
self.scheduler = optim.lr_scheduler.MultiStepLR(
self.optimizer, self.milestones, gamma=conf.train.gamma)
self.print_freq = len(self.loader) // 10
self.val_loader = get_val_loader(conf)
else:
self.model.load_state_dict(torch.load(conf.model_path))
def __init__(self, conf):
print(conf)
self.model = ResNet()
self.model.cuda()
if conf.initial:
self.model.load_state_dict(torch.load("models/"+conf.model))
print('Load model_ir_se101.pth')
self.milestones = conf.milestones
self.loader, self.class_num = get_train_loader(conf)
self.total_class = 16520
self.data_num = 285356
self.writer = SummaryWriter(conf.log_path)
self.step = 0
self.paras_only_bn, self.paras_wo_bn = separate_bn_paras(self.model)
if conf.meta:
self.head = Arcface(embedding_size=conf.embedding_size, classnum=self.total_class)
self.head.cuda()
if conf.initial:
self.head.load_state_dict(torch.load("models/head_op.pth"))
print('Load head_op.pth')
self.optimizer = RAdam([
{'params': self.paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4},
{'params': self.paras_only_bn}
], lr=conf.lr)
self.meta_optimizer = RAdam([
{'params': self.paras_wo_bn + [self.head.kernel], 'weight_decay': 5e-4},
{'params': self.paras_only_bn}
], lr=conf.lr)
self.head.train()
else:
self.head = dict()
self.optimizer = dict()
for race in races:
self.head[race] = Arcface(embedding_size=conf.embedding_size, classnum=self.class_num[race])
self.head[race].cuda()
if conf.initial:
self.head[race].load_state_dict(torch.load("models/head_op_{}.pth".format(race)))
print('Load head_op_{}.pth'.format(race))
self.optimizer[race] = RAdam([
{'params': self.paras_wo_bn + [self.head[race].kernel], 'weight_decay': 5e-4},
{'params': self.paras_only_bn}
], lr=conf.lr, betas=(0.5, 0.999))
self.head[race].train()
# self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, patience=40, verbose=True)
self.board_loss_every = min(len(self.loader[race]) for race in races) // 10
self.evaluate_every = self.data_num // 5
self.save_every = self.data_num // 2
self.eval, self.eval_issame = get_val_data(conf)
def __init__(self, conf, inference=False):
pprint.pprint(conf)
self.model = resnet18(conf.model.use_senet, conf.model.embedding_size,
conf.model.drop_out, conf.model.se_reduction)
self.model = torch.nn.DataParallel(self.model).cuda()
if not inference:
self.loader = get_train_loader(conf)
self.optimizer = optim.SGD(list(self.model.parameters()),
lr=conf.train.lr,
momentum=conf.train.momentum)
print(self.optimizer)
self.scheduler = optim.lr_scheduler.MultiStepLR(
self.optimizer, conf.train.milestones, gamma=conf.train.gamma)
print('optimizers generated')
self.print_freq = len(self.loader) // 2
print('print_freq: %d' % self.print_freq)
else:
self.test_loader = get_test_loader(conf)
import model
import data_pipe
from data_pipe import offset
from config import get_config
races = ['Caucasian', 'African', 'Asian', 'Indian']
conf = get_config()
conf.meta = True
ir_se101_model = model.ResNet()
ir_se101_model.load_state_dict(torch.load('models/model_ir_se101.pth'))
ir_se101_model.cuda()
ir_se101_model.eval()
loader, class_num = data_pipe.get_train_loader(conf)
total_class = 0
for race in races:
total_class += class_num[race]
class_count = torch.zeros(total_class)
kernel = torch.zeros(512, total_class)
for race in races:
for imgs, labels in tqdm(iter(loader[race])):
imgs = imgs.cuda()
out = ir_se101_model(imgs).cpu().detach()
for i in range(len(labels)):
label = int(labels[i])
kernel[:, label] = (class_count[label] * kernel[:, label] +
out[i]) / (class_count[label] + 1)
parser = argparse.ArgumentParser(description='Estimate Accuracy')
parser.add_argument("-r", "--race", help="target race", default='All', type=str)
parser.add_argument("-b", "--bound", help="max number of test samples for every race", default=6000, type=int)
args = parser.parse_args()
races = ['Caucasian', 'African', 'Asian', 'Indian']
if args.race == 'All':
target = races
else:
assert args.race in races
target = [args.race]
conf = get_config()
# loader, class_num = get_train_loader(conf)
loader, class_num = get_train_loader(conf)
for race in target:
ir_se101_model = model.ResNet()
# ir_se101_model.load_state_dict(torch.load('models/model_ir_se101_final_{}.pth'.format(race)))
ir_se101_model.load_state_dict(torch.load('models/model_ir_se101.pth'))
ir_se101_model.cuda()
ir_se101_model.eval()
head = model.Arcface(embedding_size=512, classnum=class_num[race])
# head.load_state_dict(torch.load("models/head_final_{}.pth".format(race)))
head.load_state_dict(torch.load("models/head_op_{}.pth".format(race)))
head.cuda()
head.eval()
correct = 0.