text: """Torch Classification Estimator"""
  # pylint: disable=unused-variable,bad-whitespace,missing-function-docstring,logging-format-interpolation,arguments-differ,logging-not-lazy, not-callable
  import math
  import copy
  import os
  import logging
  import time
  import warnings
  import pickle
  from contextlib import suppress
  from PIL import Image
  
  import pandas as pd
  import numpy as np
  import torch
  import torch.nn as nn
  import torchvision.utils
  from torch.optim.optimizer import Optimizer
  
  from timm.data import create_loader, Mixup, FastCollateMixup, AugMixDataset
  from timm.models import create_model, safe_model_name, convert_splitbn_model, model_parameters
  from timm.utils import random_seed, dispatch_clip_grad, accuracy, unwrap_model, get_state_dict
  from timm.optim import create_optimizer_v2
  from timm.utils import ApexScaler, NativeScaler, ModelEmaV2, AverageMeter
  from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
  
  from .default import TorchImageClassificationCfg
  from .utils import resolve_data_config, update_cfg, optimizer_kwargs, \
                     create_scheduler, rmse, create_optimizer_v2a
  from ..constants import BINARY, MULTICLASS, REGRESSION
  from ..utils import EarlyStopperOnPlateau
  from ..utils import _suggest_load_context
  from ..conf import _BEST_CHECKPOINT_FILE
  from ..base_estimator import BaseEstimator, set_default
  
  warnings.filterwarnings('ignore', message='.*Argument interpolation should be of type InterpolationMode instead of int.*')
  logging.getLogger('timm.models').setLevel(logging.WARNING)
  
  try:
      from apex import amp
      from apex.parallel import convert_syncbn_model
      has_apex = True
  except ImportError:
      has_apex = False
  
  has_native_amp = False
  try:
      if getattr(torch.cuda.amp, 'autocast') is not None:
          has_native_amp = True
  except AttributeError:
      pass
  
  
  @set_default(TorchImageClassificationCfg())
  class TorchImageClassificationEstimator(BaseEstimator):
      """Torch Estimator implementation for Image Classification.
  
      Parameters
      ----------
      config : dict
          Config in nested dict.
      logger : logging.Logger
          Optional logger for this estimator, can be `None` when default setting is used.
      reporter : callable
          The reporter for metric checkpointing.
      net : torch.nn.Module
          The custom network. If defined, the model name in config will be ignored so your
          custom network will be used for training rather than pulling it from model zoo.
      """
      def __init__(self, config, logger=None, reporter=None, net=None, optimizer=None, problem_type=None):
          super().__init__(config, logger=logger, reporter=reporter, name=None)
          if problem_type is None:
              problem_type = MULTICLASS
          self._problem_type = problem_type
          self._feature_net = None
          self._custom_net = False
  
          self._img_cls_cfg = self._cfg.img_cls
          self._data_cfg = self._cfg.data
          self._optimizer_cfg = self._cfg.optimizer
          self._train_cfg = self._cfg.train
          self._augmentation_cfg = self._cfg.augmentation
          self._model_ema_cfg = self._cfg.model_ema
          self._misc_cfg = self._cfg.misc
  
          # resolve AMP arguments based on PyTorch / Apex availability
          self.use_amp = None
          if self._misc_cfg.amp:
              # `amp` chooses native amp before apex (APEX ver not actively maintained)
              if self._misc_cfg.native_amp and has_native_amp:
                  self.use_amp = 'native'
              elif self._misc_cfg.apex_amp and has_apex:
                  self.use_amp = 'apex'
              elif self._misc_cfg.apex_amp or self._misc_cfg.native_amp:
                  self._logger.warning(f'Neither APEX or native Torch AMP is available, using float32. \
                                         Install NVIDA apex or upgrade to PyTorch 1.6')
          # FIXME: will provided model conflict with config provided?
          if net is not None:
              assert isinstance(net, nn.Module), f"given custom network {type(net)}, `torch.nn` expected"
              try:
                  net.to('cpu')
                  self._custom_net = True
              except ValueError:
                  pass
          self.net = net
          if optimizer is not None:
              self._logger.warning('Custom optimizer object not supported. Will follow the config instead.')
          self._optimizer = None
  
      def _fit(self, train_data, val_data, time_limit=math.inf):
          tic = time.time()
          self._cp_name = ''
          self._best_acc = -float('inf')
          self.epochs = self._train_cfg.epochs
          self.epoch = 0
          self.start_epoch = self._train_cfg.start_epoch
          self._time_elapsed = 0
          if max(self.start_epoch, self.epoch) >= self.epochs:
              return {'time', self._time_elapsed}
          self._init_trainer()
          self._init_model_ema()
          self._time_elapsed += time.time() - tic
          return self._resume_fit(train_data, val_data, time_limit=time_limit)
  
      def _resume_fit(self, train_data, val_data, time_limit=math.inf):
          tic = time.time()
          # TODO: regression not implemented
          if self._problem_type != REGRESSION and (not self.classes or not self.num_class):
              raise ValueError('This is a classification problem and we are not able to determine classes of dataset')
  
          if max(self.start_epoch, self.epoch) >= self.epochs:
              return {'time': self._time_elapsed}
  
          # wrap DP if possible
          if self.found_gpu:
              self.net = torch.nn.DataParallel(self.net, device_ids=[int(i) for i in self.valid_gpus])
          self.net = self.net.to(self.ctx[0])
  
          # prepare dataset
          train_dataset = train_data.to_torch()
          val_dataset = val_data.to_torch()
  
          # setup mixup / cutmix
          self._collate_fn = None
          self._mixup_fn = None
          self.mixup_active = self._augmentation_cfg.mixup > 0 or self._augmentation_cfg.cutmix > 0. or self._augmentation_cfg.cutmix_minmax is not None
          if self.mixup_active:
              mixup_args = dict(
                  mixup_alpha=self._augmentation_cfg.mixup, cutmix_alpha=self._augmentation_cfg.cutmix,
                  cutmix_minmax=self._augmentation_cfg.cutmix_minmax, prob=self._augmentation_cfg.mixup_prob,
                  switch_prob=self._augmentation_cfg.mixup_switch_prob, mode=self._augmentation_cfg.mixup_mode,
                  label_smoothing=self._augmentation_cfg.smoothing, num_classes=self.num_class)
              if self._misc_cfg.prefetcher:
                  self._collate_fn = FastCollateMixup(**mixup_args)
              else:
                  self._mixup_fn = Mixup(**mixup_args)
  
          # create data loaders w/ augmentation pipeiine
          train_interpolation = self._augmentation_cfg.train_interpolation
          if self._augmentation_cfg.no_aug or not train_interpolation:
              train_interpolation = self._data_cfg.interpolation
          train_loader = create_loader(
              train_dataset,
              input_size=self._data_cfg.input_size,
              batch_size=self._train_cfg.batch_size,
              is_training=True,
              use_prefetcher=self._misc_cfg.prefetcher,
              no_aug=self._augmentation_cfg.no_aug,
              scale=self._augmentation_cfg.scale,
              ratio=self._augmentation_cfg.ratio,
              hflip=self._augmentation_cfg.hflip,
              vflip=self._augmentation_cfg.vflip,
              color_jitter=self._augmentation_cfg.color_jitter,
              auto_augment=self._augmentation_cfg.auto_augment,
              interpolation=train_interpolation,
              mean=self._data_cfg.mean,
              std=self._data_cfg.std,
              num_workers=self._misc_cfg.num_workers,
              distributed=False,
              collate_fn=self._collate_fn,
              pin_memory=self._misc_cfg.pin_mem,
              use_multi_epochs_loader=self._misc_cfg.use_multi_epochs_loader
          )
  
          val_loader = create_loader(
              val_dataset,
              input_size=self._data_cfg.input_size,
              batch_size=self._data_cfg.validation_batch_size_multiplier * self._train_cfg.batch_size,
              is_training=False,
              use_prefetcher=self._misc_cfg.prefetcher,
              interpolation=self._data_cfg.interpolation,
              mean=self._data_cfg.mean,
              std=self._data_cfg.std,
              num_workers=self._misc_cfg.num_workers,
              distributed=False,
              crop_pct=self._data_cfg.crop_pct,
              pin_memory=self._misc_cfg.pin_mem,
          )
  
          self._time_elapsed += time.time() - tic
          return self._train_loop(train_loader, val_loader, time_limit=time_limit)
  
      def _train_loop(self, train_loader, val_loader, time_limit=math.inf):
          start_tic = time.time()
          # setup loss function
          if self.mixup_active:
              # smoothing is handled with mixup target transform
              train_loss_fn = SoftTargetCrossEntropy()
          elif self._augmentation_cfg.smoothing:
              train_loss_fn = LabelSmoothingCrossEntropy(smoothing=self._augmentation_cfg.smoothing)
          else:
              train_loss_fn = nn.CrossEntropyLoss()
          validate_loss_fn = nn.CrossEntropyLoss()
          train_loss_fn = train_loss_fn.to(self.ctx[0])
          validate_loss_fn = validate_loss_fn.to(self.ctx[0])
          eval_metric = self._misc_cfg.eval_metric
          if self._problem_type == REGRESSION:
              train_loss_fn = nn.MSELoss()
              validate_loss_fn = nn.MSELoss()
              eval_metric = 'rmse'
          early_stopper = EarlyStopperOnPlateau(
              patience=self._train_cfg.early_stop_patience,
              min_delta=self._train_cfg.early_stop_min_delta,
              baseline_value=self._train_cfg.early_stop_baseline,
              max_value=self._train_cfg.early_stop_max_value)
  
          self._logger.info('Start training from [Epoch %d]', max(self._train_cfg.start_epoch, self.epoch))
  
          self._time_elapsed += time.time() - start_tic
          for self.epoch in range(max(self.start_epoch, self.epoch), self.epochs):
              epoch = self.epoch
              if self._best_acc >= 1.0:
                  self._logger.info('[Epoch {}] Early stopping as acc is reaching 1.0'.format(epoch))
                  break
              should_stop, stop_message = early_stopper.get_early_stop_advice()
              if should_stop:
                  self._logger.info('[Epoch {}] '.format(epoch) + stop_message)
                  break
              train_metrics = self.train_one_epoch(
                  epoch, self.net, train_loader, self._optimizer, train_loss_fn,
                  lr_scheduler=self._lr_scheduler, output_dir=self._logdir,
                  amp_autocast=self._amp_autocast, loss_scaler=self._loss_scaler, model_ema=self._model_ema, mixup_fn=self._mixup_fn, time_limit=time_limit)
              # reaching time limit, exit early
              if train_metrics['time_limit']:
                  self._logger.warning(f'`time_limit={time_limit}` reached, exit early...')
                  return {'train_acc': train_metrics['train_acc'], 'valid_acc': self._best_acc,
                          'time': self._time_elapsed, 'checkpoint': self._cp_name}
              post_tic = time.time()
  
              eval_metrics = self.validate(self.net, val_loader, validate_loss_fn, amp_autocast=self._amp_autocast)
  
              if self._model_ema is not None and not self._model_ema_cfg.model_ema_force_cpu:
                  ema_eval_metrics = self.validate(
                      self._model_ema.module, val_loader, validate_loss_fn, amp_autocast=self._amp_autocast)
                  eval_metrics = ema_eval_metrics
  
              if self._problem_type == REGRESSION:
                  val_acc = eval_metrics['rmse']
                  if self._reporter:
                      self._reporter(epoch=epoch, acc_reward=-val_acc)
                  early_stopper.update(-val_acc)
  
                  if -val_acc > self._best_acc:
                      self._cp_name = os.path.join(self._logdir, _BEST_CHECKPOINT_FILE)
                      self._logger.info('[Epoch %d] Current best rmse: %f vs previous %f, saved to %s',
                                        self.epoch, val_acc, -self._best_acc, self._cp_name)
                      self.save(self._cp_name)
                      self._best_acc = -val_acc
              else:
                  val_acc = eval_metrics['top1']
                  if self._reporter:
                      self._reporter(epoch=epoch, acc_reward=val_acc)
                  early_stopper.update(val_acc)
  
                  if val_acc > self._best_acc:
                      self._cp_name = os.path.join(self._logdir, _BEST_CHECKPOINT_FILE)
                      self._logger.info('[Epoch %d] Current best top-1: %f vs previous %f, saved to %s',
                                        self.epoch, val_acc, self._best_acc, self._cp_name)
                      self.save(self._cp_name)
                      self._best_acc = val_acc
  
              if self._lr_scheduler is not None:
                  # step LR for next epoch
                  self._lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
  
              self._time_elapsed += time.time() - post_tic
  
          if 'accuracy' in train_metrics:
              return {'train_acc': train_metrics['accuracy'], 'valid_acc': self._best_acc,
                      'time': self._time_elapsed, 'checkpoint': self._cp_name}
          # rmse
          else:
              if self._problem_type == REGRESSION:
                  return {'train_score': train_metrics['rmse'], 'valid_score': -self._best_acc,
                          'time': self._time_elapsed, 'checkpoint': self._cp_name}
              # mixup
              else:
                  return {'train_score': train_metrics['rmse'], 'valid_acc': self._best_acc,
                          'time': self._time_elapsed, 'checkpoint': self._cp_name}
  
      def train_one_epoch(
              self, epoch, net, loader, optimizer, loss_fn,
              lr_scheduler=None, output_dir=None, amp_autocast=suppress,
              loss_scaler=None, model_ema=None, mixup_fn=None, time_limit=math.inf):
          start_tic = time.time()
          if self._augmentation_cfg.mixup_off_epoch and epoch >= self._augmentation_cfg.mixup_off_epoch:
              if self._misc_cfg.prefetcher and loader.mixup_enabled:
                  loader.mixup_enabled = False
              elif mixup_fn is not None:
                  mixup_fn.mixup_enabled = False
  
          second_order = hasattr(optimizer, 'is_second_order') and optimizer.is_second_order
          losses_m = AverageMeter()
          train_metric_score_m = AverageMeter()
  
          net.train()
  
          num_updates = epoch * len(loader)
          self._time_elapsed += time.time() - start_tic
          tic = time.time()
          last_tic = time.time()
          train_metric_name = 'accuracy'
          batch_idx = 0
          for batch_idx, (input, target) in enumerate(loader):
              b_tic = time.time()
              if self._time_elapsed > time_limit:
                  return {'train_acc': train_metric_score_m.avg, 'train_loss': losses_m.avg, 'time_limit': True}
              if self._problem_type == REGRESSION:
                  target = target.to(torch.float32)
              if not self._misc_cfg.prefetcher:
                  # prefetcher would move data to cuda by default
                  input, target = input.to(self.ctx[0]), target.to(self.ctx[0])
                  if mixup_fn is not None:
                      input, target = mixup_fn(input, target)
  
              with amp_autocast():
                  output = net(input)
                  if self._problem_type == REGRESSION:
                      output = output.flatten()
                  loss = loss_fn(output, target)
              if self._problem_type == REGRESSION:
                  train_metric_name = 'rmse'
                  train_metric_score = rmse(output, target)
              else:
                  if output.shape == target.shape:
                      train_metric_name = 'rmse'
                      train_metric_score = rmse(output, target)
                  else:
                      train_metric_score = accuracy(output, target)[0] / 100
  
              losses_m.update(loss.item(), input.size(0))
              train_metric_score_m.update(train_metric_score.item(), output.size(0))
  
              optimizer.zero_grad()
              if loss_scaler is not None:
                  loss_scaler(
                      loss, optimizer,
                      clip_grad=self._optimizer_cfg.clip_grad, clip_mode=self._optimizer_cfg.clip_mode,
                      parameters=model_parameters(net, exclude_head='agc' in self._optimizer_cfg.clip_mode),
                      create_graph=second_order)
              else:
                  loss.backward(create_graph=second_order)
                  if self._optimizer_cfg.clip_grad is not None:
                      dispatch_clip_grad(
                          model_parameters(net, exclude_head='agc' in self._optimizer_cfg.clip_mode),
                          value=self._optimizer_cfg.clip_grad, mode=self._optimizer_cfg.clip_mode)
                  optimizer.step()
  
              if model_ema is not None:
                  model_ema.update(net)
  
              if self.found_gpu:
                  torch.cuda.synchronize()
  
              num_updates += 1
              if (batch_idx+1) % self._misc_cfg.log_interval == 0:
                  lrl = [param_group['lr'] for param_group in optimizer.param_groups]
                  lr = sum(lrl) / len(lrl)
                  self._logger.info('Epoch[%d] Batch [%d]\tSpeed: %f samples/sec\t%s=%f\tlr=%f',
                                    epoch, batch_idx,
                                    self._train_cfg.batch_size*self._misc_cfg.log_interval/(time.time()-last_tic),
                                    train_metric_name, train_metric_score_m.avg, lr)
                  last_tic = time.time()
  
                  if self._misc_cfg.save_images and output_dir:
                      torchvision.utils.save_image(
                          input,
                          os.path.join(output_dir, 'train-batch-%d.jpg' % batch_idx),
                          padding=0,
                          normalize=True)
  
              if lr_scheduler is not None:
                  lr_scheduler.step_update(num_updates=num_updates, metric=losses_m.avg)
  
              self._time_elapsed += time.time() - b_tic
  
          throughput = int(self._train_cfg.batch_size * batch_idx / (time.time() - tic))
          self._logger.info('[Epoch %d] training: %s=%f', epoch, train_metric_name, train_metric_score_m.avg)
          self._logger.info('[Epoch %d] speed: %d samples/sec\ttime cost: %f', epoch, throughput, time.time()-tic)
  
          end_time = time.time()
          if hasattr(optimizer, 'sync_lookahead'):
              optimizer.sync_lookahead()
  
          self._time_elapsed += time.time() - end_time
  
          return {train_metric_name: train_metric_score_m.avg, 'train_loss': losses_m.avg, 'time_limit': False}
  
      def validate(self, net, loader, loss_fn, amp_autocast=suppress, metric_name=None):
          losses_m = AverageMeter()
          top1_m = AverageMeter()
          top5_m = AverageMeter()
          rmse_m = AverageMeter()
  
          net.eval()
  
          with torch.no_grad():
              for batch_idx, (input, target) in enumerate(loader):
                  if not self._misc_cfg.prefetcher:
                      input = input.to(self.ctx[0])
                      target = target.to(self.ctx[0])
  
                  with amp_autocast():
                      output = net(input)
                      if self._problem_type == REGRESSION:
                          output = output.flatten()
                  if isinstance(output, (tuple, list)):
                      output = output[0]
  
                  if self._problem_type == REGRESSION:
                      if metric_name:
                          assert metric_name == 'rmse', f'{metric_name} metric not supported for regression.'
                      val_metric_score = rmse(output, target)
                  else:
                      val_metric_score = accuracy(output, target, topk=(1, min(5, self.num_class)))
  
                  # augmentation reduction
                  reduce_factor = self._misc_cfg.tta
                  if self._problem_type != REGRESSION and reduce_factor > 1:
                      output = output.unfold(0, reduce_factor, reduce_factor).mean(dim=2)
                      target = target[0:target.size(0):reduce_factor]
  
                  loss = loss_fn(output, target)
                  reduced_loss = loss.data
  
                  if self.found_gpu:
                      torch.cuda.synchronize()
  
                  losses_m.update(reduced_loss.item(), input.size(0))
                  if self._problem_type == REGRESSION:
                      rmse_score = val_metric_score
                      rmse_m.update(rmse_score.item(), output.size(0))
                  else:
                      acc1, acc5 = val_metric_score
                      acc1 /= 100
                      acc5 /= 100
                      top1_m.update(acc1.item(), output.size(0))
                      top5_m.update(acc5.item(), output.size(0))
  
          if self._problem_type == REGRESSION:
              self._logger.info('[Epoch %d] validation: rmse=%f', self.epoch, rmse_m.avg)
              return {'loss': losses_m.avg, 'rmse': rmse_m.avg}
          else:
              self._logger.info('[Epoch %d] validation: top1=%f top5=%f', self.epoch, top1_m.avg, top5_m.avg)
              return {'loss': losses_m.avg, 'top1': top1_m.avg, 'top5': top5_m.avg}
  
      def _init_network(self, **kwargs):
          load_only = kwargs.get('load_only', False)
          if not self.num_class and self._problem_type != REGRESSION:
              raise ValueError('This is a classification problem and we are not able to create network when `num_class` is unknown. \
                  It should be inferred from dataset or resumed from saved states.')
          assert len(self.classes) == self.num_class
  
          # Disable syncBatchNorm as it's only supported on DDP
          if self._train_cfg.sync_bn:
              self._logger.info(
                  'Disable Sync batch norm as it is not supported for now.')
              update_cfg(self._cfg, {'train': {'sync_bn': False}})
  
          # ctx
          self.found_gpu = False
          valid_gpus = []
          if self._cfg.gpus:
              valid_gpus = self._torch_validate_gpus(self._cfg.gpus)
              self.found_gpu = True
              if not valid_gpus:
                  self.found_gpu = False
                  self._logger.warning(
                      'No gpu detected, fallback to cpu. You can ignore this warning if this is intended.')
              elif len(valid_gpus) != len(self._cfg.gpus):
                  self._logger.warning(
                      f'Loaded on gpu({valid_gpus}), different from gpu({self._cfg.gpus}).')
          self.ctx = [torch.device(f'cuda:{gid}') for gid in valid_gpus] if self.found_gpu else [torch.device('cpu')]
          self.valid_gpus = valid_gpus
  
          if not self.found_gpu and self.use_amp:
              self.use_amp = None
              self._logger.warning('Training on cpu. AMP disabled.')
              update_cfg(self._cfg, {'misc': {'amp': False, 'apex_amp': False, 'native_amp': False}})
  
          if not self.found_gpu and self._misc_cfg.prefetcher:
              self._logger.warning(
                  'Training on cpu. Prefetcher disabled.')
              update_cfg(self._cfg, {'misc': {'prefetcher': False}})
              self._logger.warning(
                  'Training on cpu. SyncBatchNorm disabled.')
              update_cfg(self._cfg, {'train': {'sync_bn': False}})
  
          random_seed(self._misc_cfg.seed)
  
          if not self.net:
              self.net = create_model(
                  self._img_cls_cfg.model,
                  pretrained=self._img_cls_cfg.pretrained and not load_only,
                  num_classes=max(self.num_class, 1),
                  global_pool=self._img_cls_cfg.global_pool_type,
                  drop_rate=self._augmentation_cfg.drop,
                  drop_path_rate=self._augmentation_cfg.drop_path,
                  drop_block_rate=self._augmentation_cfg.drop_block,
                  bn_momentum=self._train_cfg.bn_momentum,
                  bn_eps=self._train_cfg.bn_eps,
                  scriptable=self._misc_cfg.torchscript
              )
  
              self._logger.info(f'Model {safe_model_name(self._img_cls_cfg.model)} created, param count: \
                                          {sum([m.numel() for m in self.net.parameters()])}')
          else:
              self._logger.info(f'Use user provided model. Neglect model in config.')
              out_features = list(self.net.children())[-1].out_features
              if self._problem_type != REGRESSION:
                  assert out_features == self.num_class, f'Custom model out_feature {out_features} != num_class {self.num_class}.'
              else:
                  assert out_features == 1, f'Regression problem expects num_out_feature == 1, got {out_features} instead.'
  
          resolve_data_config(self._cfg, model=self.net)
  
          self.net = self.net.to(self.ctx[0])
  
          # setup synchronized BatchNorm
          if self._train_cfg.sync_bn:
              if has_apex and self.use_amp != 'native':
                  # Apex SyncBN preferred unless native amp is activated
                  self.net = convert_syncbn_model(self.net)
              else:
                  self.net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.net)
              self._logger.info(
                  'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '
                  'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.')
  
          if self._misc_cfg.torchscript:
              assert not self.use_amp == 'apex', 'Cannot use APEX AMP with torchscripted model'
              assert not self._train_cfg.sync_bn, 'Cannot use SyncBatchNorm with torchscripted model'
              self.net = torch.jit.script(self.net)
  
      def _init_trainer(self):
          if self._optimizer is None:
              if self._img_cls_cfg.pretrained and not self._custom_net \
                  and (self._train_cfg.transfer_lr_mult != 1 or self._train_cfg.output_lr_mult != 1):
                  # adjust feature/last_fc learning rate multiplier in optimizer
                  self._logger.debug(f'Reduce network lr multiplier to {self._train_cfg.transfer_lr_mult}, while keep ' +
                                     f'last FC layer lr_mult to {self._train_cfg.output_lr_mult}')
                  optim_kwargs = optimizer_kwargs(cfg=self._cfg)
                  optim_kwargs['feature_lr_mult'] = self._cfg.train.transfer_lr_mult
                  optim_kwargs['fc_lr_mult'] = self._cfg.train.output_lr_mult
                  self._optimizer = create_optimizer_v2(self.net, **optimizer_kwargs(cfg=self._cfg))
              else:
                  self._optimizer = create_optimizer_v2(self.net, **optimizer_kwargs(cfg=self._cfg))
          self._init_loss_scaler()
          self._lr_scheduler, self.epochs = create_scheduler(self._cfg, self._optimizer)
          self._lr_scheduler.step(self.start_epoch, self.epoch)
  
      def _init_loss_scaler(self):
          # setup automatic mixed-precision (AMP) loss scaling and op casting
          self._amp_autocast = suppress  # do nothing
          self._loss_scaler = None
          if self.use_amp == 'apex':
              self.net, self._optimizer = amp.initialize(self.net, self._optimizer, opt_level='O1')
              self._loss_scaler = ApexScaler()
              self._logger.info('Using NVIDIA APEX AMP. Training in mixed precision.')
          elif self.use_amp == 'native':
              self._amp_autocast = torch.cuda.amp.autocast
              self._loss_scaler = NativeScaler()
              self._logger.info('Using native Torch AMP. Training in mixed precision.')
          else:
              self._logger.info('AMP not enabled. Training in float32.')
  
      def _init_model_ema(self):
          # Disable for now
          if self._model_ema_cfg.model_ema:
              self._logger.info('Disable EMA as it is not supported for now.')
              update_cfg(self._cfg, {'model_ema': {'model_ema': False}})
          # setup exponential moving average of model weights, SWA could be used here too
          self._model_ema = None
          if self._model_ema_cfg.model_ema:
              # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
              self._model_ema = ModelEmaV2(
                  self.net, decay=self._model_ema_cfg.model_ema_decay, device='cpu' if self._model_ema_cfg.model_ema_force_cpu else None)
  
  
      def evaluate(self, val_data, metric_name=None):
          return self._evaluate(val_data, metric_name)
  
      def _evaluate(self, val_data, metric_name=None):
          if self._problem_type == REGRESSION:
              validate_loss_fn = nn.MSELoss()
          else:
              validate_loss_fn = nn.CrossEntropyLoss()
          validate_loss_fn = validate_loss_fn.to(self.ctx[0])
          val_data = val_data.to_torch()
          val_loader = create_loader(
              val_data,
              input_size=self._data_cfg.input_size,
              batch_size=self._data_cfg.validation_batch_size_multiplier * self._train_cfg.batch_size,
              is_training=False,
              use_prefetcher=self._misc_cfg.prefetcher,
              interpolation=self._data_cfg.interpolation,
              mean=self._data_cfg.mean,
              std=self._data_cfg.std,
              num_workers=self._misc_cfg.num_workers,
              distributed=False,
              crop_pct=self._data_cfg.crop_pct,
              pin_memory=self._misc_cfg.pin_mem,
          )
          return self.validate(self.net, val_loader, validate_loss_fn, amp_autocast=self._amp_autocast, metric_name=metric_name)
  
      def _predict(self, x, **kwargs):
          with_proba = kwargs.get('with_proba', False)
          if with_proba and self._problem_type not in [MULTICLASS, BINARY]:
              raise AssertionError('with_proba is only supported for classification problems. Please use predict instead.')
          if isinstance(x, str):
              return self._predict((x,), **kwargs).drop(columns=['image'], errors='ignore')
          elif isinstance(x, pd.DataFrame):
              assert 'image' in x.columns, "Expect column `image` for input images"
              df = self._predict(tuple(x['image']), **kwargs)
              return df.reset_index(drop=True)
          elif isinstance(x, (list, tuple)):
              loader = create_loader(
                  ImageListDataset(x),
                  input_size=self._data_cfg.input_size,
                  batch_size=self._train_cfg.batch_size,
                  use_prefetcher=self._misc_cfg.prefetcher,
                  interpolation=self._data_cfg.interpolation,
                  mean=self._data_cfg.mean,
                  std=self._data_cfg.std,
                  num_workers=self._misc_cfg.num_workers,
                  crop_pct=self._data_cfg.crop_pct
              )
  
              self.net.eval()
  
              topk = min(5, self.num_class)
              results = []
              idx = 0
              with torch.no_grad():
                  for input, _ in loader:
                      input = input.to(self.ctx[0])
                      labels = self.net(input)
                      for l in labels:
                          if self._problem_type in [MULTICLASS, BINARY]:
                              probs = nn.functional.softmax(l, dim=0).cpu().numpy().flatten()
                              if with_proba:
                                  results.append({'image_proba': probs.tolist(), 'image': x[idx]})
                              else:
                                  topk_inds = l.topk(topk)[1].cpu().numpy().flatten()
                                  results.extend([{'class': self.classes[topk_inds[k]],
                                                   'score': probs[topk_inds[k]],
                                                   'id': topk_inds[k],
                                                   'image': x[idx]}
                                                  for k in range(topk)])
                          else:
                              results.append({'prediction': l.cpu().numpy().flatten(), 'image': x[idx]})
                          idx += 1
              return pd.DataFrame(results)
          elif not isinstance(x, torch.Tensor):
              raise ValueError('Input is not supported: {}'.format(type(x)))
          with torch.no_grad():
              input = x.to(self.ctx[0])
              label = self.net(input)
              if self._problem_type in [MULTICLASS, BINARY]:
                  topk = min(5, self.num_class)
                  probs = nn.functional.softmax(label, dim=0).cpu().numpy().flatten()
                  topk_inds = label.topk(topk)[1].cpu().numpy().flatten()
                  if with_proba:
                      df = pd.DataFrame([{'image_proba': probs.tolist()}])
                  else:
                      df = pd.DataFrame([{'class': self.classes[topk_inds[k]],
                                          'score': probs[topk_inds[k]],
                                          'id': topk_inds[k]}
                                         for k in range(topk)])
              else:
                  df = pd.DataFrame([{'prediction': label.cpu().numpy().flatten()}])
          return df
  
  
      def _predict_feature(self, x, **kwargs):
          if isinstance(x, str):
              return self._predict_feature((x,))
          elif isinstance(x, pd.DataFrame):
              assert 'image' in x.columns, "Expect column `image` for input images"
              df = self._predict_feature(tuple(x['image']))
              df = df.set_index(x.index)
              df['image'] = x['image']
              return df
          elif isinstance(x, (list, tuple)):
              assert isinstance(x[0], str), "expect image paths in list/tuple input"
              loader = create_loader(
                  ImageListDataset(x),
                  input_size=self._data_cfg.input_size,
                  batch_size=self._train_cfg.batch_size,
                  use_prefetcher=self._misc_cfg.prefetcher,
                  interpolation=self._data_cfg.interpolation,
                  mean=self._data_cfg.mean,
                  std=self._data_cfg.std,
                  num_workers=self._misc_cfg.num_workers,
                  crop_pct=self._data_cfg.crop_pct
              )
  
              self.net.eval()
  
              results = []
              with torch.no_grad():
                  for input, _ in loader:
                      input = input.to(self.ctx[0])
                      try:
                          features = self.net.forward_features(input)
                      except AttributeError:
                          features = self.net.module.forward_features(input)
                      for f in features:
                          f = f.cpu().numpy().flatten()
                          results.append({'image_feature': f})
              df = pd.DataFrame(results)
              df['image'] = x
              return df
          elif not isinstance(x, torch.Tensor):
              raise ValueError('Input is not supported: {}'.format(type(x)))
          with torch.no_grad():
              input = x.to(self.ctx[0])
              feature = self.net.forward_features(input)
              result = [{'image_feature': feature}]
          df = pd.DataFrame(result)
          return df
  
      def _reconstruct_state_dict(self, state_dict):
          new_state_dict = {}
          for k, v in state_dict.items():
              name = k[7:] if k.startswith('module') else k
              new_state_dict[name] = v
          return new_state_dict
  
      def save(self, filename):
          d = dict()
          current_states = self.__dict__.copy()
          if self.net:
              if not self._custom_net:
                  if isinstance(self.net, torch.nn.DataParallel):
                      d['model_state_dict'] = get_state_dict(self.net.module, unwrap_model)
                  else:
                      d['model_state_dict'] = get_state_dict(self.net, unwrap_model)
              else:
                  net_pickle = pickle.dumps(self.net)
                  d['net_pickle'] = net_pickle
              self.net = None
          if self._optimizer:
              d['optimizer_state_dict'] = self._optimizer.state_dict()
              self._optimizer = None
          if hasattr(self, '_loss_scaler') and self._loss_scaler:
              d[self._loss_scaler.state_dict_key] = self._loss_scaler.state_dict()
              d['_loss_scaler_state_dict_key'] = self._loss_scaler.state_dict_key
          if self._model_ema:
              d['ema_state_dict'] = get_state_dict(self._model_ema, unwrap_model)
              self._model_ema = None
          self._logger = None
          self._reporter = None
          d['estimator'] = self
          torch.save(d, filename)
          self.__dict__.update(current_states)
  
      @classmethod
      def load(cls, filename, ctx='auto'):
          d = torch.load(filename, map_location=torch.device('cpu'))
          est = d.pop('estimator')
          # logger
          est._logger = logging.getLogger(cls.__name__)
          est._logger.setLevel(logging.ERROR)
          try:
              fh = logging.FileHandler(est._log_file)
              est._logger.addHandler(fh)
          #pylint: disable=bare-except
          except:
              pass
          model_state_dict = d.get('model_state_dict', None)
          net_pickle = d.get('net_pickle', None)
          if model_state_dict:
              est._init_network(load_only=True)
              net_state_dict = est._reconstruct_state_dict(model_state_dict)
              if isinstance(est.net, torch.nn.DataParallel):
                  est.net.module.load_state_dict(net_state_dict)
              else:
                  est.net.load_state_dict(net_state_dict)
          elif net_pickle:
              est.net = pickle.loads(net_pickle)
          optimizer_state_dict = d.get('optimizer_state_dict', None)
          if optimizer_state_dict:
              est._init_trainer()
              est._optimizer.load_state_dict(optimizer_state_dict)
          if hasattr(est, '_loss_scaler') and est._loss_scaler:
              loss_scaler_state_dict_key = d.get('loss_scaler_state_dict')
              loss_scaler_dict = d.get(loss_scaler_state_dict_key, None)
              if loss_scaler_dict:
                  est._loss_scaler.load_state_dict(loss_scaler_dict)
          ema_state_dict = d.get('ema_state_dict', None)
          est._init_model_ema()
          if ema_state_dict:
              ema_state_dict = est._reconstruct_state_dict(ema_state_dict)
              if isinstance(est.net, torch.nn.DataParallel):
                  est._model_ema.module.module.load_state_dict(ema_state_dict)
              else:
                  est._model_ema.module.load_state_dict(ema_state_dict)
          new_ctx = _suggest_load_context(est.net, ctx, est.ctx)
          est.reset_ctx(new_ctx)
          est._logger.setLevel(logging.INFO)
          return est
  
      # pylint: disable=redefined-outer-name, reimported
      def __getstate__(self):
          d = self.__dict__.copy()
          try:
              import torch
              net = d.pop('net', None)
              model_ema = d.pop('_model_ema', None)
              optimizer = d.pop('_optimizer', None)
              loss_scaler = d.pop('_loss_scaler', None)
              save_state = {}
              if net is not None:
                  if not self._custom_net:
                      if isinstance(net, torch.nn.DataParallel):
                          save_state['state_dict'] = get_state_dict(net.module, unwrap_model)
                      else:
                          save_state['state_dict'] = get_state_dict(net, unwrap_model)
                  else:
                      net_pickle = pickle.dumps(net)
                      save_state['net_pickle'] = net_pickle
              if optimizer is not None:
                  save_state['optimizer'] = optimizer.state_dict()
              if loss_scaler is not None:
                  save_state[loss_scaler.state_dict_key] = loss_scaler.state_dict()
              if model_ema is not None:
                  save_state['state_dict_ema'] = get_state_dict(model_ema, unwrap_model)
          except ImportError:
              pass
          d['save_state'] = save_state
          d['_logger'] = None
          d['_reporter'] = None
          return d
  
      def __setstate__(self, state):
          save_state = state.pop('save_state', None)
          self.__dict__.update(state)
          # logger
          self._logger = logging.getLogger(state.get('_name', self.__class__.__name__))
          self._logger.setLevel(logging.ERROR)
          try:
              fh = logging.FileHandler(self._log_file)
              self._logger.addHandler(fh)
          #pylint: disable=bare-except
          except:
              pass
          if not save_state:
              self.net = None
              self._optimizer = None
              self._logger.setLevel(logging.INFO)
              return
          try:
              import torch
              self.net = None
              self._optimizer = None
              if self._custom_net:
                  if save_state.get('net_pickle', None):
                      self.net = pickle.loads(save_state['net_pickle'])
              else:
                  if save_state.get('state_dict', None):
                      self._init_network(load_only=True)
                      net_state_dict = self._reconstruct_state_dict(save_state['state_dict'])
                      if isinstance(self.net, torch.nn.DataParallel):
                          self.net.module.load_state_dict(net_state_dict)
                      else:
                          self.net.load_state_dict(net_state_dict)
              if save_state.get('optimizer', None):
                  self._init_trainer()
                  self._optimizer.load_state_dict(save_state['optimizer'])
              if hasattr(self, '_loss_scaler') and self._loss_scaler and self._loss_scaler.state_dict_key in save_state:
                  loss_scaler_dict = save_state[self._loss_scaler.state_dict_key]
                  self._loss_scaler.load_state_dict(loss_scaler_dict)
              if save_state.get('state_dict_ema', None):
                  self._init_model_ema()
                  model_ema_dict = save_state.get('state_dict_ema')
                  model_ema_dict = self._reconstruct_state_dict(model_ema_dict)
                  if isinstance(self.net, torch.nn.DataParallel):
                      self._model_ema.module.module.load_state_dict(model_ema_dict)
                  else:
                      self._model_ema.module.load_state_dict(model_ema_dict)
          except ImportError:
              pass
          self._logger.setLevel(logging.INFO)
  
  
  class ImageListDataset(torch.utils.data.Dataset):
      """An internal image list dataset for batch predict"""
      def __init__(self, imlist):
          self._imlist = imlist
          self.transform = None
  
      def __getitem__(self, idx):
          img = Image.open(self._imlist[idx]).convert('RGB')
          label = None
          if self.transform is not None:
              img = self.transform(img)
          return img, torch.tensor(-1, dtype=torch.long)
  
      def __len__(self):
          return len(self._imlist)
  