Training

import os

import pytorch_lightning as pl
import torch

import fastface as ff

# set seed
pl.seed_everything(41)

# build training transforms
train_transforms = ff.transforms.Compose(
    ff.transforms.Interpolate(target_size=480),
    ff.transforms.Padding(target_size=(480, 480)),
    ff.transforms.RandomHorizontalFlip(p=0.5),
)

# build val transforms
val_transforms = ff.transforms.Compose(
    ff.transforms.Interpolate(target_size=480),
    ff.transforms.Padding(target_size=(480, 480)),
)

# build torch.utils.data.DataLoader for training
train_dl = ff.dataset.FDDBDataset(
    phase="train", transforms=train_transforms
).get_dataloader(batch_size=8, shuffle=True, num_workers=8)

# build torch.utils.data.DataLoader for validation
val_dl = ff.dataset.FDDBDataset(phase="val", transforms=val_transforms).get_dataloader(
    batch_size=8, shuffle=False, num_workers=8
)

# define preprocess dict
preprocess = {"mean": 127.5, "std": 127.5, "normalized_input": False}

# define hyper parameter dict
hparams = {
    "learning_rate": 0.1,
    "momentum": 0.9,
    "weight_decay": 0.00001,
    "milestones": [500000, 1000000, 1500000],
    "gamma": 0.1,
    "ratio": 10,
}

# checkout available architectures to train
print(ff.list_archs())
# ["lffd"]
arch = "lffd"

# checkout available configs for the architecture
print(ff.list_arch_configs(arch))
# ["original", "slim"]
config = "slim"

# build pl.LightningModule with random weights
model = ff.FaceDetector.build(
    arch, config=config, preprocess=preprocess, hparams=hparams
)

# add average precision pl.metrics.Metric to the model
model.add_metric("average_precision", ff.metric.AveragePrecision(iou_threshold=0.5))

model_save_name = "{}_{}_{}_best".format(arch, config, "fddb")
ckpt_save_path = "./checkpoints"

# resume with checkpoint, if exists
ckpt_resume_path = os.path.join(ckpt_save_path, model_save_name + ".ckpt")
if not os.path.isfile(ckpt_resume_path):
    ckpt_resume_path = None

# define checkpoint callback
checkpoint_callback = pl.callbacks.ModelCheckpoint(
    dirpath=ckpt_save_path,
    verbose=True,
    filename=model_save_name,
    monitor="average_precision",
    save_top_k=1,
    mode="max",  # only pick max of `average_precision`
)

# define pl.Trainer
trainer = pl.Trainer(
    default_root_dir=".",
    accumulate_grad_batches=4,  # update weights every 4 batches
    callbacks=[checkpoint_callback],
    gpus=1 if torch.cuda.is_available() else 0,
    precision=32,
    resume_from_checkpoint=ckpt_resume_path,
    max_epochs=100,
    check_val_every_n_epoch=2,  # run validation every 2 epochs
    gradient_clip_val=10,  # clip gradients
)

# start training
trainer.fit(model, train_dataloader=train_dl, val_dataloaders=[val_dl])