# torchreid.engine

## Base Engine

class `torchreid.engine.engine.``Engine`(datamanager, model, optimizer=None, scheduler=None, use_cpu=False)

  A generic base Engine class for both image- and video-reid.

  Parameters:

  - datamanager (DataManager) – an instance of `torchreid.data.ImageDataManager` or `torchreid.data.VideoDataManager`.
  - model (nn.Module) – model instance.
  - optimizer (Optimizer) – an Optimizer.
  - scheduler (LRScheduler, optional) – if None, no learning rate decay will be performed.
  - use_cpu (bool, optional) – use cpu. Default is False.

  `run`(save_dir='log', max_epoch=0, start_epoch=0, fixbase_epoch=0, open_layers=None, start_eval=0, eval_freq=-1, test_only=False, print_freq=10, dist_metric='euclidean', visrank=False, visrank_topk=20, use_metric_cuhk03=False, ranks=[1, 5, 10, 20])

    A unified pipeline for training and evaluating a model.

    Parameters:

    - save_dir (str) – directory to save model.
    - max_epoch (int) – maximum epoch.
    - start_epoch (int, optional) – starting epoch. Default is 0.
    - fixbase_epoch (int, optional) – number of epochs to train `open_layers` (new layers) while keeping base layers frozen. Default is 0. `fixbase_epoch` is not counted in `max_epoch`.
    - open_layers (str or list, optional) – layers (attribute names) open for training.
    - start_eval (int, optional) – from which epoch to start evaluation. Default is 0.
    - eval_freq (int, optional) – evaluation frequency. Default is -1 (meaning evaluation is only performed at the end of training).
    - test_only (bool, optional) – if True, only runs evaluation on test datasets. Default is False.
    - print_freq (int, optional) – print_frequency. Default is 10.
    - dist_metric (str, optional) – distance metric used to compute distance matrix between query and gallery. Default is “euclidean”.
    - visrank (bool, optional) – visualizes ranked results. Default is False. Visualization will be performed every test time, so it is recommended to enable `visrank` when `test_only` is True. The ranked images will be saved to “save_dir/ranks-epoch/dataset_name”, e.g. “save_dir/ranks-60/market1501”.
    - visrank_topk (int, optional) – top-k ranked images to be visualized. Default is 20.
    - use_metric_cuhk03 (bool, optional) – use single-gallery-shot setting for cuhk03. Default is False. This should be enabled when using cuhk03 classic split.
    - ranks (list, optional) – cmc ranks to be computed. Default is [1, 5, 10, 20].

  `test`(epoch, testloader, dist_metric='euclidean', visrank=False, visrank_topk=20, save_dir='', use_metric_cuhk03=False, ranks=[1, 5, 10, 20])

    Tests model on target datasets.

    Note

    This function has been called in `run()` when necessary.

    Note

    The test pipeline implemented in this function suits both image- and video-reid. In general, a subclass of Engine only needs to re-implement `_extract_features()` and `_parse_data_for_eval()` when necessary, but not a must. Please refer to the source code for more details.

    Parameters:

    - epoch (int) – current epoch.
    - testloader (dict) – dictionary containing {dataset_name: ‘query’: queryloader, ‘gallery’: galleryloader}.
    - dist_metric (str, optional) – distance metric used to compute distance matrix between query and gallery. Default is “euclidean”.
    - visrank (bool, optional) – visualizes ranked results. Default is False. Visualization will be performed every test time, so it is recommended to enable `visrank` when `test_only` is True. The ranked images will be saved to “save_dir/ranks-epoch/dataset_name”, e.g. “save_dir/ranks-60/market1501”.
    - visrank_topk (int, optional) – top-k ranked images to be visualized. Default is 20.
    - save_dir (str) – directory to save visualized results if `visrank` is True.
    - use_metric_cuhk03 (bool, optional) – use single-gallery-shot setting for cuhk03. Default is False. This should be enabled when using cuhk03 classic split.
    - ranks (list, optional) – cmc ranks to be computed. Default is [1, 5, 10, 20].

  `train`()

    Performs training on source datasets for one epoch.

    This will be called every epoch in `run()`, e.g.

        for epoch in range(start_epoch, max_epoch):
            self.train(some_arguments)

    Note

    This needs to be implemented in subclasses.

## Image Engines

class `torchreid.engine.image.softmax.``ImageSoftmaxEngine`(datamanager, model, optimizer, scheduler=None, use_cpu=False, label_smooth=True)

  Softmax-loss engine for image-reid.

  Parameters:

  - datamanager (DataManager) – an instance of `torchreid.data.ImageDataManager` or `torchreid.data.VideoDataManager`.
  - model (nn.Module) – model instance.
  - optimizer (Optimizer) – an Optimizer.
  - scheduler (LRScheduler, optional) – if None, no learning rate decay will be performed.
  - use_cpu (bool, optional) – use cpu. Default is False.
  - label_smooth (bool, optional) – use label smoothing regularizer. Default is True.

  Examples:

      import torch
      import torchreid
      datamanager = torchreid.data.ImageDataManager(
          root='path/to/reid-data',
          sources='market1501',
          height=256,
          width=128,
          combineall=False,
          batch_size=32
      )
      model = torchreid.models.build_model(
          name='resnet50',
          num_classes=datamanager.num_train_pids,
          loss='softmax'
      )
      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
      )
      engine = torchreid.engine.ImageSoftmaxEngine(
          datamanager, model, optimizer, scheduler=scheduler
      )
      engine.run(
          max_epoch=60,
          save_dir='log/resnet50-softmax-market1501',
          print_freq=10
      )

  `train`(epoch, trainloader, fixbase=False, open_layers=None, print_freq=10)

    Trains the model for one epoch on source datasets using softmax loss.

    Parameters:

    - epoch (int) – current epoch.
    - trainloader (Dataloader) – training dataloader.
    - fixbase (bool, optional) – whether to fix base layers. Default is False.
    - open_layers (str or list, optional) – layers open for training.
    - print_freq (int, optional) – print frequency. Default is 10.

class `torchreid.engine.image.triplet.``ImageTripletEngine`(datamanager, model, optimizer, margin=0.3, weight_t=1, weight_x=1, scheduler=None, use_cpu=False, label_smooth=True)

  Triplet-loss engine for image-reid.

  Parameters:

  - datamanager (DataManager) – an instance of `torchreid.data.ImageDataManager` or `torchreid.data.VideoDataManager`.
  - model (nn.Module) – model instance.
  - optimizer (Optimizer) – an Optimizer.
  - margin (float, optional) – margin for triplet loss. Default is 0.3.
  - weight_t (float, optional) – weight for triplet loss. Default is 1.
  - weight_x (float, optional) – weight for softmax loss. Default is 1.
  - scheduler (LRScheduler, optional) – if None, no learning rate decay will be performed.
  - use_cpu (bool, optional) – use cpu. Default is False.
  - label_smooth (bool, optional) – use label smoothing regularizer. Default is True.

  Examples:

      import torch
      import torchreid
      datamanager = torchreid.data.ImageDataManager(
          root='path/to/reid-data',
          sources='market1501',
          height=256,
          width=128,
          combineall=False,
          batch_size=32,
          num_instances=4,
          train_sampler='RandomIdentitySampler' # this is important
      )
      model = torchreid.models.build_model(
          name='resnet50',
          num_classes=datamanager.num_train_pids,
          loss='triplet'
      )
      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
      )
      engine = torchreid.engine.ImageTripletEngine(
          datamanager, model, optimizer, margin=0.3,
          weight_t=0.7, weight_x=1, scheduler=scheduler
      )
      engine.run(
          max_epoch=60,
          save_dir='log/resnet50-triplet-market1501',
          print_freq=10
      )

  `train`(epoch, trainloader, fixbase=False, open_layers=None, print_freq=10)

    Trains the model for one epoch on source datasets using hard mining triplet loss.

    Parameters:

    - epoch (int) – current epoch.
    - trainloader (Dataloader) – training dataloader.
    - fixbase (bool, optional) – whether to fix base layers. Default is False.
    - open_layers (str or list, optional) – layers open for training.
    - print_freq (int, optional) – print frequency. Default is 10.

## Video Engines

class `torchreid.engine.video.softmax.``VideoSoftmaxEngine`(datamanager, model, optimizer, scheduler=None, use_cpu=False, label_smooth=True, pooling_method='avg')

  Softmax-loss engine for video-reid.

  Parameters:

  - datamanager (DataManager) – an instance of `torchreid.data.ImageDataManager` or `torchreid.data.VideoDataManager`.
  - model (nn.Module) – model instance.
  - optimizer (Optimizer) – an Optimizer.
  - scheduler (LRScheduler, optional) – if None, no learning rate decay will be performed.
  - use_cpu (bool, optional) – use cpu. Default is False.
  - label_smooth (bool, optional) – use label smoothing regularizer. Default is True.
  - pooling_method (str, optional) – how to pool features for a tracklet. Default is “avg” (average). Choices are [“avg”, “max”].

  Examples:

      import torch
      import torchreid
      # Each batch contains batch_size*seq_len images
      datamanager = torchreid.data.VideoDataManager(
          root='path/to/reid-data',
          sources='mars',
          height=256,
          width=128,
          combineall=False,
          batch_size=8, # number of tracklets
          seq_len=15 # number of images in each tracklet
      )
      model = torchreid.models.build_model(
          name='resnet50',
          num_classes=datamanager.num_train_pids,
          loss='softmax'
      )
      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
      )
      engine = torchreid.engine.VideoSoftmaxEngine(
          datamanager, model, optimizer, scheduler=scheduler,
          pooling_method='avg'
      )
      engine.run(
          max_epoch=60,
          save_dir='log/resnet50-softmax-mars',
          print_freq=10
      )

class `torchreid.engine.video.triplet.``VideoTripletEngine`(datamanager, model, optimizer, margin=0.3, weight_t=1, weight_x=1, scheduler=None, use_cpu=False, label_smooth=True, pooling_method='avg')

  Triplet-loss engine for video-reid.

  Parameters:

  - datamanager (DataManager) – an instance of `torchreid.data.ImageDataManager` or `torchreid.data.VideoDataManager`.
  - model (nn.Module) – model instance.
  - optimizer (Optimizer) – an Optimizer.
  - margin (float, optional) – margin for triplet loss. Default is 0.3.
  - weight_t (float, optional) – weight for triplet loss. Default is 1.
  - weight_x (float, optional) – weight for softmax loss. Default is 1.
  - scheduler (LRScheduler, optional) – if None, no learning rate decay will be performed.
  - use_cpu (bool, optional) – use cpu. Default is False.
  - label_smooth (bool, optional) – use label smoothing regularizer. Default is True.
  - pooling_method (str, optional) – how to pool features for a tracklet. Default is “avg” (average). Choices are [“avg”, “max”].

  Examples:

      import torch
      import torchreid
      # Each batch contains batch_size*seq_len images
      # Each identity is sampled with num_instances tracklets
      datamanager = torchreid.data.VideoDataManager(
          root='path/to/reid-data',
          sources='mars',
          height=256,
          width=128,
          combineall=False,
          num_instances=4,
          train_sampler='RandomIdentitySampler'
          batch_size=8, # number of tracklets
          seq_len=15 # number of images in each tracklet
      )
      model = torchreid.models.build_model(
          name='resnet50',
          num_classes=datamanager.num_train_pids,
          loss='triplet'
      )
      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
      )
      engine = torchreid.engine.VideoTripletEngine(
          datamanager, model, optimizer, margin=0.3,
          weight_t=0.7, weight_x=1, scheduler=scheduler,
          pooling_method='avg'
      )
      engine.run(
          max_epoch=60,
          save_dir='log/resnet50-triplet-mars',
          print_freq=10
      )