Laughing-q/MetricTrainer

MetricTrainer

• This framework is for metric study, built on timm and pytorch-metric-learning.
• This repo is experimental.

TODO

• logger
• eval
• DDP
• amp training
• dataloader
• saving model stuff
• resume
• plot results
• support other optimizer
• support other scheduler
• ema
• inference

Quick Start

Installation

Clone repo and install requirements.txt in a Python>=3.7.0 environment, includingPyTorch>=1.7.1.

pip install pytorch-metric-learning
pip install git+https://github.com/Laughing-q/pytorch-image-models.git
git clone https://github.com/Laughing-q/MetricTrainer.git
cd MetricTrainer
pip install -r requirements.txt
pip install -e .

Training

• Prepare your own config, see configs/ for more details.
• partial fc rely on DDP mode, so if you get only one GPU, just set n=1.
• Multi GPU(DDP)

python -m torch.distributed.run --nproc_per_node=n tools/train.py -c configs/test.yaml

Eval

• prepare your val datasets like below(you can get the datasets from insightface):

├── root_dir
│   ├── lfw.bin
│   ├── cfp_fp.bin
│   ├── agedb_30.bin
│   ├── calfw.bin
│   ├── cplfw.bin
│   ├── vgg2_fp.bin

python tools/eval.py -c configs/partial_glint360k.yaml -d root_dir \
    -w runs/CosFace_noaug/best.pt

Tips

• I haven't test resume at all, maybe there will get some bug.