text: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import copy
  import os
  import time
  import numpy as np
  import torch
  import torch.optim as optim
  import torch.nn.functional as F
  import torch.nn as nn
  from torchvision import datasets, transforms
  import util.dataset
  import util.io
  
  
  class Patience:
      def __init__(self, capacity = 10):
          self._tolerance = self._capacity = int(capacity)
          self._calmness = 0.0
      
      def __call__(self):
          if self._tolerance > 0:
              return True
          else:
              return False
  
      def test(self, calmness):
          if calmness > self._calmness:
              self._calmness = calmness
              self._tolerance = self._capacity
          else:
              self._tolerance -= 1
          return self
  
  def train(args, model, data_loaders, device, writer):
      torch.manual_seed(args.seed)
  
      train_loader, validation_loader, test_loader = data_loaders
  
      if hasattr(model, 'criterion'):
          criterion = getattr(model, 'criterion')
      else:
          print("No criterion in model. Assuming nn.CrossEntropyLoss().")
          criterion = nn.CrossEntropyLoss()
  
      if hasattr(model, 'optimizer'):
          optimizer = model.optimizer
      else:
          #optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
          optimizer = optim.Adam(model.parameters(), lr=args.lr)
      print(optimizer)
      if hasattr(model, 'scheduler'):
          print(model.scheduler)
  
      model.to(device)
      if args.load_last:
          print("Loading saved model: {}".format(args.load_last_file))
          util.io.load(args.load_last_file, model, optimizer)
          state = 0, 0.0, 0, float("Inf") # best_epoch, best_accuracy, best_loss_epoch, best_loss
          print("Testing loaded model...")
          accuracy, state = test_epoch(0, args, model, test_loader, criterion, optimizer, device, writer, 'test/', state)
  
      state = 0, 0.0, 0, float("Inf") # best_epoch, best_accuracy, best_loss_epoch, best_loss
      train_state = state
      patient = Patience(args.patience)
      epoch = 1
      while patient() or args.patience == 0:
          if hasattr(model, 'scheduler'):
              model.scheduler.step()
          if args.reverse_layer_binarization:
              if hasattr(model, 'binarize_layer_reverse'):
                  model.binarize_layer_reverse()
              else:
                  raise NotImplementedError("reverse_layer_binarization requested, but binarize_layer_reverse() does not exist")
          else:
              if hasattr(model, 'binarize_layer'):
                  model.binarize_layer()
          train_epoch(epoch, args, model, train_loader, criterion, optimizer, device, writer)
          if args.test_validation:
              accuracy, state = test_epoch(epoch, args, model, validation_loader, criterion, optimizer, device, writer, 'val/', state)
          else:
              accuracy, state = test_epoch(epoch, args, model, test_loader, criterion, optimizer, device, writer, 'test/', state)
  
          if args.test_training:
              _, train_state = test_epoch(epoch, args, model, train_loader, criterion, optimizer, device, writer, 'train/', train_state)
  
          if hasattr(model, 'scheduler'):
              writer.add_scalar('train/learning_rate', model.scheduler.get_last_lr()[0], epoch)
  
          if args.patience == 0:
              if epoch >= args.epochs: break
          else:
              patient.test(accuracy)
          epoch += 1
  
  
      best_epoch, best_accuracy, best_loss_epoch, best_loss = state
      print('Best accuracy {:.4f}% at epoch {}'.format(100*best_accuracy, best_epoch))
      print('Best loss {:.4g} at epoch {}'.format(best_loss, best_loss_epoch))
      if args.test_training:
          best_epoch, best_accuracy, best_loss_epoch, best_loss = train_state
          print('Best training accuracy {:.4f}% at epoch {}'.format(100*best_accuracy, best_epoch))
          print('Best training loss {:.4g} at epoch {}'.format(best_loss, best_loss_epoch))
      if args.save_last:
          util.io.save(epoch, model, optimizer, best_loss, args.logdir, pattern='last-{}.tar')
  
  
  def train_epoch(epoch, args, model, data_loader, criterion, optimizer, device, writer):
      print("Logdir {}".format(args.logdir))
      print("Model {} {}".format(args.model, model.__class__.__name__))
      model.train()
      pid = os.getpid()
      start_epoch = time.time()
      for batch_idx, (data, target) in enumerate(data_loader, 0):
          start = time.time()
          data, target = data.to(device), target.to(device)
          optimizer.zero_grad()
          #weight_stats("{} {:>3} optimizer.zero_grad()".format(epoch, batch_idx), model)
          output = model(data)
          loss = criterion(output, target)
          loss.backward()
          #weight_stats("{} {:>3} loss.backward()      ".format(epoch, batch_idx), model)
          if hasattr(model, 'restore_full_precision'):
              #print("=== before")
              #model.report_weight_stats()
              model.restore_full_precision()
              #print("=== after")
              #model.report_weight_stats()
          optimizer.step()
          if hasattr(model, 'clip_weights'):
              model.clip_weights()
          #weight_stats("{} {:>3} optimizer.step()     ".format(epoch, batch_idx), model)
          if batch_idx % args.log_interval == 0:
              print('{}  Train Epoch: {} [{:>5}/{:>5} {:.0f}%]  Loss: {:.6g}'.format(
                  pid, epoch, batch_idx * len(data), len(data_loader.dataset),
                  100. * batch_idx / len(data_loader), loss.item()))
      print('Epoch duration: {:.4f} seconds'.format(time.time() - start_epoch))
      record_weight_histogram(args, epoch, model, writer)
  
  
  def record_weight_histogram(args, epoch, model, writer):
      if args.record_histogram:
          start = time.time()
          for name, param in model.named_parameters():
              x = param.clone().cpu().data.numpy()
              writer.add_histogram(name+"/train", x, epoch)
              mantissa, exponent = np.frexp(x)
              writer.add_histogram(name+'.mantissa/train', mantissa, epoch)
              writer.add_histogram(name+'.exponent/train', exponent, epoch)
          print('Weight histogram written in {:.4f} seconds'.format(time.time() - start))
  
  
  def test_epoch(epoch, args, model, data_loader, criterion, optimizer, device, writer, writer_prefix, state):
      model.eval()
      loss = 0.0
      correct = 0
      start = time.time()
      #criterion = getattr(model, 'criterion')
      with torch.no_grad():
          if hasattr(model, 'restore_binary_weight'):
              print("test_epoch::restore_binary_weight")
              model.restore_binary_weight()
          if hasattr(model, 'report_weight_stats'):
              print("test_epoch::report_weight_stats")
              maximum, minimum, num_zeros, num_binary = model.report_weight_stats()
              writer.add_scalar('{}weight/max'.format(writer_prefix), maximum, epoch)
              writer.add_scalar('{}weight/min'.format(writer_prefix), minimum, epoch)
              writer.add_scalar('{}weight/num_zeros'.format(writer_prefix), num_zeros, epoch)
              writer.add_scalar('{}weight/num_binary'.format(writer_prefix), num_binary, epoch)
          for data, target in data_loader:
              data, target = data.to(device), target.to(device)
              output = model(data)
              #loss += F.nll_loss(output, target).item() # sum up batch loss
              loss += criterion(output, target).item() # sum up batch loss
              pred = output.max(1)[1] # get the index of the max log-probability
              correct += pred.eq(target).sum().item()
          if hasattr(model, 'restore_full_precision'):
              print("test_epoch::restore_full_precision")
              model.restore_full_precision()
  
      loss /= len(data_loader.dataset)
      accuracy = correct / len(data_loader.dataset)
      print('Prefix {}, Average loss: {:.6g}, Accuracy: {}/{} ({:.4f}%), Duration: {:.4f} seconds'.format(
          writer_prefix,
          loss, 
          correct, len(data_loader.dataset), 100. * accuracy, 
          time.time() - start))
  
      writer.add_scalar('{}accuracy'.format(writer_prefix), accuracy, epoch)
      writer.add_scalar('{}error'.format(writer_prefix), 1.0-accuracy, epoch)
      writer.add_scalar('{}loss'.format(writer_prefix), loss, epoch)
  
      best_epoch, best_accuracy, best_loss_epoch, best_loss = state
      check = epoch, accuracy
      best = best_epoch, best_accuracy
      prefix = "accuracy_"+writer_prefix.replace("/", "")
      best_epoch, best_accuracy = save_progress(prefix, args, model, optimizer, check, best)
      check = epoch, loss
      best = best_loss_epoch, best_loss
      prefix = "loss_"+writer_prefix.replace("/", "")
      best_loss_epoch, best_loss = save_progress(prefix, args, model, optimizer, check, best, higher_better=False)
      
      writer.add_scalar('{}best_accuracy'.format(writer_prefix), best_accuracy, epoch)
      writer.add_scalar('{}best_error'.format(writer_prefix), 1.0 - best_accuracy, epoch)
      writer.add_scalar('{}best_loss'.format(writer_prefix), best_loss, epoch)
      #writer.add_scalar('{}learning_rate'.format(writer_prefix), scheduler.get_lr(), epoch)
  
      state = best_epoch, best_accuracy, best_loss_epoch, best_loss
      return accuracy, state
  
  
  def save_progress(name, args, model, optimizer, check, best, higher_better=True):
      epoch, metric = check
      best_epoch, best_metric = best
      if higher_better:
          better = metric > best_metric
      else:
          better = metric < best_metric
      if better:
          print("Better {} {:.4f} at epoch {}".format(name, metric, epoch))
          if args.save_progress:
              util.io.save(epoch, model, optimizer, metric, args.logdir, pattern=name+'-{}.tar')
          return epoch, metric
      else:
          return best_epoch, best_metric
  
  
  def isnan(tensor, message):
      if torch.isnan(tensor):
          print(message)
  
  
  def weight_stats(prefix, model):
      with torch.no_grad():
          for name, param in model.named_parameters():
              high = torch.max(param)
              low = torch.min(param)
              print("{} {:>.6g} {:>.6g} {} {}".format(prefix, high, low, name, param.size()))
  