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)