def create_mtcnn_net(self, p_model_path=None, r_model_path=None, o_model_path=None, use_cuda=True):
dirname, _ = os.path.split(p_model_path)
checkpoint = CheckPoint(dirname)
pnet, rnet, onet = None, None, None
self.device = torch.device(
"cuda:0" if use_cuda and torch.cuda.is_available() else "cpu")
if p_model_path is not None:
pnet = PNet()
pnet_model_state = checkpoint.load_model(p_model_path)
pnet = checkpoint.load_state(pnet, pnet_model_state)
if (use_cuda):
pnet.to(self.device)
pnet.eval()
if r_model_path is not None:
rnet = RNet()
rnet_model_state = checkpoint.load_model(r_model_path)
rnet = checkpoint.load_state(rnet, rnet_model_state)
if (use_cuda):
rnet.to(self.device)
rnet.eval()
if o_model_path is not None:
onet = ONet()
onet_model_state = checkpoint.load_model(o_model_path)
onet = checkpoint.load_state(onet, onet_model_state)
if (use_cuda):
onet.to(self.device)
onet.eval()
return pnet, rnet, onet
'pin_memory': True
} if use_cuda else {}
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
train_loader = torch.utils.data.DataLoader(FaceDataset(config.annoPath,
transform=transform,
is_train=True),
batch_size=config.batchSize,
shuffle=True,
**kwargs)
# Set model
model = RNet()
model = model.to(device)
# Set checkpoint
checkpoint = CheckPoint(config.save_path)
# Set optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=config.step,
gamma=0.1)
# Set trainer
logger = Logger(config.save_path)
trainer = RNetTrainer(config.lr, train_loader, model, optimizer, scheduler,
logger, device)
torch.cuda.manual_seed(config.manualSeed)
device = torch.device("cuda:0" if use_cuda else "cpu")
torch.backends.cudnn.benchmark = True
# Set dataloader
kwargs = {'num_workers': config.nThreads, 'pin_memory': True} if use_cuda else {}
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
train_loader = torch.utils.data.DataLoader(
FaceDataset(config.annoPath, transform=transform, is_train=True), batch_size=config.batchSize, shuffle=True, **kwargs)
# Set model
model = RNet()
model = model.to(device)
# Set checkpoint
checkpoint = CheckPoint(config.save_path)
# Set optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=config.step, gamma=0.1)
# Set trainer
logger = Logger(config.save_path)
trainer = RNetTrainer(config.lr, train_loader, model, optimizer, scheduler, logger, device)
for epoch in range(1, config.nEpochs + 1):
cls_loss_, box_offset_loss, total_loss, accuracy = trainer.train(epoch)
checkpoint.save_model(model, index=epoch)