Torchreid Documentation

This example demonstrates the standard workflow for training a person re-identification model: importing the library, initializing an ImageDataManager, building a model, optimizer, and scheduler, and finally running the training loop using ImageSoftmaxEngine.

import torchreid

# 1. Load data manager
datamanager = torchreid.data.ImageDataManager(
    root="reid-data",
    sources="market1501",
    targets="market1501",
    height=256,
    width=128,
    batch_size_train=32,
    batch_size_test=100,
    transforms=["random_flip", "random_crop"]
)

# 2. Build model, optimizer and lr_scheduler
model = torchreid.models.build_model(
    name="resnet50",
    num_classes=datamanager.num_train_pids,
    loss="softmax",
    pretrained=True
)

model = model.cuda()

optimizer = torchreid.optim.build_optimizer(
    model,
    optim="adam",
    lr=0.0003
)

scheduler = torchreid.optim.build_lr_scheduler(
    optimizer,
    lr_scheduler="single_step",
    stepsize=20
)

# 3. Build engine
engine = torchreid.engine.ImageSoftmaxEngine(
    datamanager,
    model,
    optimizer=optimizer,
    scheduler=scheduler,
    label_smooth=True
)

# 4. Run training and test
engine.run(
    save_dir="log/resnet50",
    max_epoch=60,
    eval_freq=10,
    print_freq=10,
    test_only=False
)

To perform the neural architecture search (NAS) to find the optimal OSNet+IN design, run the main.py script using the nas.yaml configuration file. Ensure your re-ID data is stored in a directory and provided via the --root flag.

Important Notes:

• Hardware: The default configuration is optimized for 8 Tesla V100 32GB GPUs. You may need to adjust the batch size in the config file to match your available device memory.
• Evaluation: Do not use the test results obtained at the end of the architecture search to judge model performance. The results are not meaningful due to stochastic sampling layers used during the search process.
• Final Model: To evaluate the discovered architecture, you must manually construct the found architecture in osnet_child.py, then perform a standard re-train and evaluation on your re-ID datasets.

python main.py --config-file nas.yaml --root $DATA

To use a custom dataset in Torchreid, follow these three steps:

1. Create a dataset class: Inherit from torchreid.data.ImageDataset (for images) or torchreid.data.VideoDataset (for video). You must generate three lists: train, query, and gallery. Each list must contain tuples of (img_path, pid, camid):
   • img_path (str): Absolute path to the image.
   • pid (int): 0-based person ID.
   • camid (int): 0-based camera ID.
   • Note: query and gallery must share the same pid scope. train, query, and gallery must share the same camid scope.
2. Register the dataset: Use torchreid.data.register_image_dataset to make your class available via a string name.
3. Initialize an ImageDataManager: Use the registered name in the sources argument.

from torchreid.data import ImageDataset
import os.path as osp

class NewDataset(ImageDataset):
    dataset_dir = 'new_dataset'

    def __init__(self, root='', **kwargs):
        self.root = osp.abspath(osp.expanduser(root))
        self.dataset_dir = osp.join(self.root, self.dataset_dir)

        # Lists must contain (img_path, pid, camid)
        train = [('/path/to/img1.jpg', 0, 0), ...]
        query = [('/path/to/img2.jpg', 0, 1), ...]
        gallery = [('/path/to/img3.jpg', 0, 1), ...]

        super(NewDataset, self).__init__(train, query, gallery, **kwargs)

import torchreid
torchreid.data.register_image_dataset('new_dataset', NewDataset)

# Initialize Data Manager
datamanager = torchreid.data.ImageDataManager(
    root='reid-data',
    sources='new_dataset'
)