With the provided CLI, you can train several CNNs, from custom to pretrained, on the UTKFace dataset. The dataset is available here. Besides the relatively small UTK Face Dataset, it is also possible to use the much larger B3F-Dataset (B3FD), which is around 6GB in size.
Before using the CLI, you need to install the package. Refer to the Installation Instructions for more information.
Then, you can use the CLI to train a model:
train <option>=<value> <option>=<value> ...
e.g. to train the EfficientNetV2B0-Model on the UTKFace dataset with default parameters, run:
train model=efficientnetv2 model.version=B0
By default it uses the UTKFace dataset. You can also train on B3FD, using dataset=b3fd
as additional configuration.
Also note, that by default it logs to Weights & Biases. You can disable this by setting wandb.mode=offline
.
Run train --help
to see a quick overview for the most important options. This will also list all default values.
If you run train
without any options, it will use the default values for all options which are defined in src/conf/
. You can override these values by passing them as arguments to the CLI.
- dataset: either
utkface
orb3fd
- lr_schedule: either
cosine_decay
orexponential_decay
- model: many options available, e.g.
efficientnetv2
orbaselinecnn
. Seetrain --help
for all options - optimizer: either
adam
oradamw
Specify the number of folds, e.g. train.cv_folds=5
to train with 5-Fold Cross Validation. This will train K
models and evaluate them on the test set. The final score test score is calculated on a hold-out test set (different from the test set within the Cross Validation) using the ensemble prediction (average the output of all K
models).
If train.cv_folds=null
(default) it will use a ordinary train-test-split, as specified with train.test_size
(default: 0.2
).
This test size determines the size of the hold-out test set, which is used to calculate the final test score in cross validation.
For most pretrained models, you can specify the version/size of the model seperately.
E.g. for efficientnetv2
you can specify model.version=B0
to use the B0 version of the model. See here for a detailed specification of the pretrained models on Keras.
The pretrained models are trained in two-stage fashion as it is described on Keras here.
- Train only the top layers of the model with a larger learning rate.
- Fine-tune the whole (or a part of the) model with a smaller learning rate.
This behavior is determined based on the settings of
model.freeze_base
, model.finetune_base
and model.num_finetune_layers
.
freeze_base
: Iftrue
, it will freeze the base model (default:true
). This means that only the top layers are trained.finetune_base
: Iftrue
, it will fine-tune the base model. This means that the whole model is trained in the second stage. This setting depends on the model in use.num_finetune_layers
: The number of layers to fine-tune. Ifall
, it will fine-tune the whole (base) model. If notnull
it must be an integer greater 0. E.g. usemodel.num_finetune_layers=20
to finetune the top 20 layers of the base model (ignores the top layers outside of the base model)
.. warning:: Using an integer to specify the number of layers to fine-tune does not work currently. Use all
instead.
You can specify the top layer architecture with
model.top_layer_architecture._target_
. This must point to valid reference in src.top_layer_architectures
For example: The resnet top is specified via model.top_layer_architecture._target_='src.top_layer_architectures.resnet_top'
Available top layer architectures are
src.top_layer_architectures.resnet_top
(usually performs well)src.top_layer_architectures.pass_through
src.top_layer_architectures.vgg_top
src.top_layer_architectures.fully_connected
src.top_layer_architectures.fully_connected_with_dropout
src.top_layer_architectures.conv_with_fc_top
(not recommended)
This is an overview of other important options to configure the training process.
train.target_size
: Specify the size of the images (default:[150, 150]
). Note that some models require a minimum size. ForVGGFace
this must be at least[224, 224]
. ForSENet50
it must be at least[197, 197]
.train.cache_dataset
: IfTrue
, it will cache the dataset in memory. This is recommended for the UTKFace dataset, but not for B3FD, because it is too large.train.batch_size
: The batch size for training. Default:32
train.epochs
: The number of epochs to train. Default:150
lr_schedule.stage_1.initial_learning_rate
: The initial learning rate for the first stage of training (andstage_2
analogously)
augment.active
: If True
, it will use data augmentation. Default: True
. You can specify the augmentation options in augment.factors
. E.g. augment.factors.random_rotation=0.1
to specify the rotation range or augment.factors.random_rotation=null
to disable only the rotation.
Available augmentations:
random_rotation
: Randomly rotate the image by a given anglerandom_brightness
: Randomly change the brightness of the imagerandom_translation
: Randomly translate th imagerandom_flip
: Randomly flip the image horizontally or vertically (or both)
callbacks.early_stopping_patience
: Set the early stopping patience (default:5
)callbacks.model_ckpt
: Iftrue
, it will save the best model checkpoint after each epoch. Default:false
callbacks.with_wandb_ckpt
: Iftrue
, it will save the best model checkpoint to Weights & Biases. Default:false
.callbacks.visualize_predictions
: Iftrue
, it will visualize the predictions of the model on the validation set after each epoch. Default:false
. This requires wandb to be active.
wandb.mode
: eitheronline
oroffline
. Ifonline
, it will log to Weights & Biases. Ifoffline
, it will still use Weights & Biases, but it does not require an account.wandb.project
: The name of the Weights & Biases project. Default:UTKFace-v2
- Create a virtual environment with Python 3.9 or higher. It is highly recommended to use Conda because of the Tensorflow dependency.
- Install the package with
pip install .
orpoetry install
This is already provided via the Makefile
. You can just run
make setup
to create a conda environment and install the package into it.
- Install Poetry
- Set up the environment:
make setup
To install new PyPI packages, run:
poetry add <package-name>
To add dev-dependencies, run:
poetry add <package-name> --group dev
The Documentation is automatically deployed to GitHub Pages.
To view the documentation locally, run:
make docs_view
This project was generated with the Light-weight Python Template by Moritz Mistol.