Source code for torchreid.engine.image.softmax

from __future__ import division, print_function, absolute_import

from torchreid import metrics
from torchreid.losses import CrossEntropyLoss

from ..engine import Engine


[docs]class ImageSoftmaxEngine(Engine): r"""Softmax-loss engine for image-reid. Args: 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_gpu (bool, optional): use gpu. Default is True. label_smooth (bool, optional): use label smoothing regularizer. Default is True. Examples:: 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 ) """ def __init__( self, datamanager, model, optimizer, scheduler=None, use_gpu=True, label_smooth=True ): super(ImageSoftmaxEngine, self).__init__(datamanager, use_gpu) self.model = model self.optimizer = optimizer self.scheduler = scheduler self.register_model('model', model, optimizer, scheduler) self.criterion = CrossEntropyLoss( num_classes=self.datamanager.num_train_pids, use_gpu=self.use_gpu, label_smooth=label_smooth ) def forward_backward(self, data): imgs, pids = self.parse_data_for_train(data) if self.use_gpu: imgs = imgs.cuda() pids = pids.cuda() outputs = self.model(imgs) loss = self.compute_loss(self.criterion, outputs, pids) self.optimizer.zero_grad() loss.backward() self.optimizer.step() loss_summary = { 'loss': loss.item(), 'acc': metrics.accuracy(outputs, pids)[0].item() } return loss_summary