def main(args):
MODEL_DICT = {
'PFLD': PFLD,
'PFLD_Ghost': PFLD_Ghost,
'PFLD_Ghost_Slim': PFLD_Ghost_Slim,
}
MODEL_TYPE = args.model_type
WIDTH_FACTOR = args.width_factor
INPUT_SIZE = args.input_size
LANDMARK_NUMBER = args.landmark_number
model = MODEL_DICT[MODEL_TYPE](WIDTH_FACTOR, INPUT_SIZE,
LANDMARK_NUMBER).to(args.device)
checkpoint = torch.load(args.model_path, map_location=args.device)
model.load_state_dict(checkpoint)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
wlfw_val_dataset = WLFWDatasets(args.test_dataset, transform)
wlfw_val_dataloader = DataLoader(wlfw_val_dataset,
batch_size=1,
shuffle=False,
num_workers=8)
validate(model, wlfw_val_dataloader, args)
def _data_loader(self):
transform = torchvision.transforms.Compose(
[torchvision.transforms.ToTensor()])
self.data['train_loader'] = DataLoader(
WLFWDatasets(self.args.train_file, transform),
batch_size=self.args.train_batchsize,
shuffle=True,
num_workers=self.args.workers,
drop_last=False)
self.data['eval_loader'] = DataLoader(
WLFWDatasets(self.args.eval_file, transform),
batch_size=self.args.val_batchsize,
shuffle=False,
num_workers=self.args.workers)
print('Data loading was finished ...')
def main(args):
checkpoint = torch.load(args.model_path, map_location=device)
plfd_backbone = PFLDInference(args.r).to(device)
plfd_backbone = nn.DataParallel(plfd_backbone)
plfd_backbone.load_state_dict(checkpoint['plfd_backbone'])
transform = transforms.Compose([transforms.ToTensor()])
wlfw_val_dataset = WLFWDatasets(args.test_dataset, transform)
wlfw_val_dataloader = DataLoader(wlfw_val_dataset, batch_size=1, shuffle=False, num_workers=0)
validate(wlfw_val_dataloader, plfd_backbone)
def main(args):
checkpoint = torch.load(args.model_path)
plfd_backbone = PFLDInference().cuda()
auxiliarynet = AuxiliaryNet().cuda()
plfd_backbone.load_state_dict(checkpoint['plfd_backbone'])
auxiliarynet.load_state_dict(checkpoint['auxiliarynet'])
transform = transforms.Compose([transforms.ToTensor()])
wlfw_val_dataset = WLFWDatasets(args.test_dataset, transform)
wlfw_val_dataloader = DataLoader(
wlfw_val_dataset, batch_size=8, shuffle=False, num_workers=0)
validate(wlfw_val_dataloader, plfd_backbone, auxiliarynet)
def main(args):
# Step 1: parse args config
logging.basicConfig(
format=
'[%(asctime)s] [p%(process)s] [%(pathname)s:%(lineno)d] [%(levelname)s] %(message)s',
level=logging.INFO,
handlers=[
logging.FileHandler(args.log_file, mode='w'),
logging.StreamHandler()
])
print_args(args)
# Step 2: model, criterion, optimizer, scheduler
plfd_backbone = PFLDInference().cuda()
auxiliarynet = AuxiliaryNet().cuda()
criterion = PFLDLoss()
optimizer = torch.optim.Adam([{
'params': plfd_backbone.parameters()
}, {
'params': auxiliarynet.parameters()
}],
lr=args.base_lr,
weight_decay=args.weight_decay)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode='min', patience=args.lr_patience, verbose=True)
# step 3: data
# argumetion
transform = transforms.Compose([transforms.ToTensor()])
wlfwdataset = WLFWDatasets(args.dataroot, transform)
dataloader = DataLoader(wlfwdataset,
batch_size=args.train_batchsize,
shuffle=True,
num_workers=args.workers,
drop_last=False)
wlfw_val_dataset = WLFWDatasets(args.val_dataroot, transform)
wlfw_val_dataloader = DataLoader(wlfw_val_dataset,
batch_size=args.val_batchsize,
shuffle=False,
num_workers=args.workers)
# step 4: run
writer = SummaryWriter(args.tensorboard)
for epoch in range(args.start_epoch, args.end_epoch + 1):
weighted_train_loss, train_loss = train(dataloader, plfd_backbone,
auxiliarynet, criterion,
optimizer, epoch)
filename = os.path.join(str(args.snapshot),
"checkpoint_epoch_" + str(epoch) + '.pth.tar')
save_checkpoint(
{
'epoch': epoch,
'plfd_backbone': plfd_backbone.state_dict(),
'auxiliarynet': auxiliarynet.state_dict()
}, filename)
val_loss = validate(wlfw_val_dataloader, plfd_backbone, auxiliarynet,
criterion, epoch)
scheduler.step(val_loss)
writer.add_scalar('data/weighted_loss', weighted_train_loss, epoch)
writer.add_scalars('data/loss', {
'val loss': val_loss,
'train loss': train_loss
}, epoch)
writer.close()
def main(args):
# Step 1: parse args config
logging.basicConfig(
format=
'[%(asctime)s] [p%(process)s] [%(pathname)s:%(lineno)d] [%(levelname)s] %(message)s',
level=logging.INFO,
handlers=[
logging.FileHandler(args.log_file, mode='w'),
logging.StreamHandler()
])
print_args(args)
# Step 2: model, criterion, optimizer, scheduler
if wandb.config.pfld_backbone == "GhostNet":
plfd_backbone = CustomizedGhostNet(width=wandb.config.ghostnet_width, dropout=0.2)
logger.info(f"Using GHOSTNET with width={wandb.config.ghostnet_width} as backbone of PFLD backbone")
# If using pretrained weight from ghostnet model trained on image net
if (wandb.config.ghostnet_with_pretrained_weight_image_net == True):
logger.info(f"Using pretrained weights of ghostnet model trained on image net data ")
plfd_backbone = load_pretrained_weight_imagenet_for_ghostnet_backbone(
plfd_backbone, "./checkpoint_imagenet/state_dict_93.98.pth")
else:
plfd_backbone = PFLDInference().to(device) # MobileNet2 defaut
logger.info("Using MobileNet2 as backbone of PFLD backbone")
auxiliarynet = AuxiliaryNet().to(device)
# Watch model by wandb
wandb.watch(plfd_backbone)
wandb.watch(auxiliarynet)
criterion = PFLDLoss()
optimizer = torch.optim.Adam(
[{
'params': plfd_backbone.parameters()
}, {
'params': auxiliarynet.parameters()
}],
lr=args.base_lr,
weight_decay=args.weight_decay)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=args.lr_patience, verbose=True)
# step 3: data
# argumetion
transform = transforms.Compose([transforms.ToTensor()])
wlfwdataset = WLFWDatasets(args.dataroot, transform)
dataloader = DataLoader(
wlfwdataset,
batch_size=args.train_batchsize,
shuffle=True,
num_workers=args.workers,
drop_last=False)
wlfw_val_dataset = WLFWDatasets(args.val_dataroot, transform)
wlfw_val_dataloader = DataLoader(
wlfw_val_dataset,
batch_size=args.val_batchsize,
shuffle=False,
num_workers=args.workers)
# step 4: run
writer = SummaryWriter(args.tensorboard)
for epoch in range(args.start_epoch, args.end_epoch + 1):
weighted_train_loss, train_loss = train(dataloader, plfd_backbone, auxiliarynet,
criterion, optimizer, epoch)
filename = os.path.join(
str(args.snapshot), "checkpoint_epoch_" + str(epoch) + '.pth.tar')
save_checkpoint({
'epoch': epoch,
'plfd_backbone': plfd_backbone.state_dict(),
'auxiliarynet': auxiliarynet.state_dict()
}, filename)
val_loss = validate(wlfw_val_dataloader, plfd_backbone, auxiliarynet,
criterion)
wandb.log({"metric/val_loss": val_loss})
scheduler.step(val_loss)
writer.add_scalar('data/weighted_loss', weighted_train_loss, epoch)
writer.add_scalars('data/loss', {'val loss': val_loss, 'train loss': train_loss}, epoch)
writer.close()