# torchreid.engine.image.softmax — torchreid 1.4.0 documentation

# 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

    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