trainer.py

# encoding: utf-8
"""
@author:  zhoumi
@contact: zhoumi281571814@126.com
"""
import math
import time
import numpy as np

class AverageMeter(object):
    def __init__(self):
        self.n = 0
        self.sum = 0.0
        self.var = 0.0
        self.val = 0.0
        self.mean = np.nan
        self.std = np.nan

    def update(self, value, n=1):
        self.val = value
        self.sum += value
        self.var += value * value
        self.n += n

        if self.n == 0:
            self.mean, self.std = np.nan, np.nan
        elif self.n == 1:
            self.mean, self.std = self.sum, np.inf
        else:
            self.mean = self.sum / self.n
            self.std = math.sqrt(
                (self.var - self.n * self.mean * self.mean) / (self.n - 1.0))

    def value(self):
        return self.mean, self.std

    def reset(self):
        self.n = 0
        self.sum = 0.0
        self.var = 0.0
        self.val = 0.0
        self.mean = np.nan
        self.std = np.nan

class cls_tripletTrainer:
    def __init__(self, opt, model, optimzier, criterion, summary_writer):
        self.opt = opt
        self.model = model
        self.optimizer= optimzier
        self.criterion = criterion
        self.summary_writer = summary_writer

    def train(self, epoch, data_loader):
        self.model.train()

        batch_time = AverageMeter()
        data_time = AverageMeter()
        losses = AverageMeter()

        start = time.time()
        for i, inputs in enumerate(data_loader):
            data_time.update(time.time() - start)

            # model optimizer
            self._parse_data(inputs)
            self._forward()
            self.optimizer.zero_grad()
            self._backward()
            self.optimizer.step()

            batch_time.update(time.time() - start)
            losses.update(self.loss.item())

            # tensorboard
            global_step = epoch * len(data_loader) + i
            self.summary_writer.add_scalar('loss', self.loss.item(), global_step)
            self.summary_writer.add_scalar('lr', self.optimizer.param_groups[0]['lr'], global_step)

            start = time.time()

            if (i + 1) % self.opt.print_freq == 0:
                print('Epoch: [{}][{}/{}]\t'
                      'Batch Time {:.3f} ({:.3f})\t'
                      'Data Time {:.3f} ({:.3f})\t'
                      'Loss {:.3f} ({:.3f})\t'
                      .format(epoch, i + 1, len(data_loader),
                              batch_time.val, batch_time.mean,
                              data_time.val, data_time.mean,
                              losses.val, losses.mean))
        param_group = self.optimizer.param_groups
        print('Epoch: [{}]\tEpoch Time {:.3f} s\tLoss {:.3f}\t'
              'Lr {:.2e}'
              .format(epoch, batch_time.sum, losses.mean, param_group[0]['lr']))
        print()

    def _parse_data(self, inputs):
        imgs, pids, _ = inputs
        self.data = imgs.cuda()
        self.target = pids.cuda()

    def _forward(self):
        score, feat = self.model(self.data)
        if 'pcb' == self.opt.model_name:
            loss = 0
            for i in range(self.opt.num_parts):
                loss += self.criterion(score[i], feat, self.target)
            self.loss = loss / self.opt.num_parts
        elif 'MGN' == self.opt.model_name or 'StackPCBv2' == self.opt.model_name\
                or 'drop_block' == self.opt.model_name or 'StackPCB' == self.opt.model_name:
            tri_loss = [self.criterion(score[0], f, self.target, weight_s = 0) for f in feat]
            soft_loss = [self.criterion(s, feat[0], self.target, weight_t = 0) for s in score]

            self.loss = sum(tri_loss) / len(feat) + sum(soft_loss) / len(score)

        else:
            self.loss = self.criterion(score, feat, self.target)

    def _backward(self):
        self.loss.backward()