docstr: Train one epoch.
docstr_tok: ["Train","one","epoch","."]
code: def _train_epoch(self, epoch, optimizer, arch_train=False):
      """
      Train one epoch.
      """
      # switch to train mode
      self.model.train()
      self.auxiliarynet.train()
  
      batch_time = AverageMeter("batch_time")
      data_time = AverageMeter("data_time")
      losses = AverageMeter("losses")
  
      data_loader = (
          self.valid_loader if arch_train else self.train_loader
      )
      end = time.time()
      for i, (img, landmark_gt, angle_gt) in enumerate(data_loader):
          data_time.update(time.time() - end)
          img = img.to(self.device, non_blocking=True)
          landmark_gt = landmark_gt.to(self.device, non_blocking=True)
          angle_gt = angle_gt.to(self.device, non_blocking=True)
  
          lands, feats = self.model(img)
          landmarks = lands.squeeze()
          angle = self.auxiliarynet(feats)
  
          # task loss
          weighted_loss, l2_loss = self.criterion(
              landmark_gt, angle_gt, angle, landmarks
          )
          loss = l2_loss if arch_train else weighted_loss
  
          # hardware-aware loss
          perf_cost = self._get_perf_cost(requires_grad=True)
          regu_loss = self.reg_loss(perf_cost)
          if self.mode.startswith("mul"):
              loss = loss * regu_loss
          elif self.mode.startswith("add"):
              loss = loss + regu_loss
  
          # compute gradient and do SGD step
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
          # measure elapsed time
          batch_time.update(time.time() - end)
          end = time.time()
          # measure accuracy and record loss
          losses.update(
              np.squeeze(loss.cpu().detach().numpy()), img.size(0)
          )
  
          if i % 10 == 0:
              batch_log = (
                  "Train [{0}][{1}]\t"
                  "Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
                  "Data {data_time.val:.3f} ({data_time.avg:.3f})\t"
                  "Loss {losses.val:.4f} ({losses.avg:.4f})".format(
                      epoch + 1,
                      i,
                      batch_time=batch_time,
                      data_time=data_time,
                      losses=losses,
                  )
              )
              self.logger.info(batch_log)
  
file_path: nni/examples/nas/legacy/oneshot/pfld/lib/trainer.py
imports: from __future__ import absolute_import, division, print_function
  import os
  import time
  import torch
  import numpy
  from nni.algorithms.nas.pytorch.fbnet import FBNetTrainer
  from nni.nas.pytorch.utils import AverageMeter
  from .utils import accuracy
label_desc: Use image node_id 3 for calling the PFLDTrainer obj's underlying member method code with example usage: obj._train_epoch(epoch, optimizer, arch_train) without return types