text: import argparse
  import os
  import json
  import numpy as np
  import time
  from data import AttributeDataset
  from models import LinearProjection, MLP, Base
  from models import ProxyNet, ProxyLoss
  from utils import AverageMeter, get_data_stats
  from validate import calibrate
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  from torch.optim.lr_scheduler import CosineAnnealingLR
  
  
  # Training settings
  parser = argparse.ArgumentParser(description='PyTorch DML')
  # hyper-parameters
  parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                      help='input batch size for training (default: 64)')
  parser.add_argument('--epochs', type=int, default=30, metavar='N',
                      help='number of epochs to train (default: 30)')
  parser.add_argument('--start_epoch', type=int, default=1, metavar='N',
                      help='number of start epoch (default: 1)')
  parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                      help='learning rate (default: 1e-2)')
  parser.add_argument('--seed', type=int, default=456, metavar='S',
                      help='random seed (default: 1)')
  # flags
  parser.add_argument('--no-cuda', action='store_true', default=False,
                      help='de-enables CUDA training')
  parser.add_argument('--no-decay', action='store_true', default=False,
                      help='de-enables LR decay')
  # model
  parser.add_argument('--dim-input', type=int, default=2048, metavar='N',
                      help='input dimensions (default: 2048)')
  parser.add_argument('--temp', type=float, default=1, metavar='LR',
                      help='softmax temperature (default: 1)')
  parser.add_argument('--ent', type=float, default=0.1, metavar='LR',
                      help='amount of entropy reg (default: 0.1)')
  parser.add_argument('--margin', type=float, default=0.2, metavar='LR',
                      help='margin in the entropy reg (default: 0.2)')
  parser.add_argument('--validation', action='store_true', default=False,
                      help='Use validation sets, does not work with gan path')
  parser.add_argument('--nhidden', type=int, default=1024,
                      help='number of hidden units')
  parser.add_argument('--mlp', action='store_true', default=False,
                      help='MLP instead of linear projection (default: False)')
  parser.add_argument('--normalize', action='store_true', default=False,
                      help='normalize data between 0 and 1')
  # plumbing
  parser.add_argument('--dataset', type=str, required=True,
                      help='sun|cub|awa1')
  parser.add_argument('--sentences', action='store_true', default=False,
                      help='Use sentences, only for CUB')
  parser.add_argument('--gan-path', type=str, default=None,
                      help='Path to GAN features')
  parser.add_argument('--data-dir', type=str, metavar='DD', default='data',
                      help='data folder path')
  parser.add_argument('--exp-dir', type=str, default='exp', metavar='ED',
                      help='folder for saving exp')
  parser.add_argument('--m', type=str, default=None, metavar='M',
                      help='message')
  
  
  args = parser.parse_args()
  args.cuda = not args.no_cuda and torch.cuda.is_available()
  args.lr_decay = not args.no_decay
  
  torch.manual_seed(args.seed)
  if args.cuda:
      torch.cuda.manual_seed(args.seed)
  
  # basic stats
  if args.validation:
      n_seen, n_val_unseen, n_unseen, n_attr = get_data_stats(
          args.dataset, args.validation)
  else:
      n_seen, n_unseen, n_attr = get_data_stats(args.dataset)
  args.n_classes = n_seen
  
  
  if args.sentences:
      args.dim_embed = 1024
  else:
      args.dim_embed = n_attr
  
  # create experiment folder
  basic = '%s' % (args.dataset)
  if args.m is not None:
      basic = basic + '_' + args.m
  path = os.path.join(args.exp_dir, basic)
  if not os.path.exists(path):
      os.makedirs(path)
  
  # saving logs
  with open(os.path.join(path, 'config.json'), 'w') as f:
      json.dump(args.__dict__, f, indent=4, sort_keys=True)
  
  
  def write_logs(txt, logpath=os.path.join(path, 'logs.txt')):
      with open(logpath, 'a') as f:
          f.write('\n')
          f.write(txt)
  
  
  def main():
      if args.gan_path is None:
          both = False
      else:
          both = True
  
      if args.validation:
          train_loader = torch.utils.data.DataLoader(
              AttributeDataset(args.data_dir, args.dataset,
                               features_path=args.gan_path,
                               mode='train',
                               both=both,
                               normalize=args.normalize,
                               sentences=args.sentences),
              batch_size=args.batch_size,
              shuffle=True)
          val_seen_loader = torch.utils.data.DataLoader(
              AttributeDataset(args.data_dir, args.dataset,
                               features_path=args.gan_path,
                               mode='val_seen',
                               generalized=True,
                               normalize=args.normalize,
                               sentences=args.sentences),
              batch_size=args.batch_size,
              shuffle=False)
          val_unseen_loader = torch.utils.data.DataLoader(
              AttributeDataset(args.data_dir, args.dataset,
                               features_path=args.gan_path,
                               mode='val_unseen',
                               generalized=True,
                               normalize=args.normalize,
                               sentences=args.sentences),
              batch_size=args.batch_size,
              shuffle=False)
      else:
          trainval_loader = torch.utils.data.DataLoader(
              AttributeDataset(args.data_dir, args.dataset,
                               features_path=args.gan_path,
                               mode='trainval',
                               both=both,
                               normalize=args.normalize,
                               sentences=args.sentences),
              batch_size=args.batch_size,
              shuffle=True)
  
      test_seen_loader = torch.utils.data.DataLoader(
          AttributeDataset(args.data_dir, args.dataset,
                           features_path=args.gan_path,
                           mode='test_seen',
                           generalized=True,
                           normalize=args.normalize,
                           sentences=args.sentences),
          batch_size=args.batch_size,
          shuffle=False)
  
      test_unseen_loader = torch.utils.data.DataLoader(
          AttributeDataset(args.data_dir, args.dataset,
                           features_path=args.gan_path,
                           mode='test_unseen',
                           generalized=True,
                           normalize=args.normalize,
                           sentences=args.sentences),
          batch_size=args.batch_size,
          shuffle=False)
  
      # instanciate the models
      if args.mlp:
          mlp = MLP(args.dim_input, [args.nhidden*2], args.nhidden)
      else:
          mlp = LinearProjection(args.dim_input, args.nhidden)
      embed = LinearProjection(args.nhidden, args.dim_embed)
  
      if args.sentences:
          cam_key = 'sentences'
      else:
          cam_key = 'emb'
  
      if args.validation:
          cam = torch.from_numpy(train_loader.dataset.data[cam_key].T)
      else:
          cam = torch.from_numpy(trainval_loader.dataset.data[cam_key].T)
      proxies = ProxyNet(args.n_classes, args.dim_embed, proxies=cam)
  
      model = Base(mlp, embed, proxies)
  
      criterion = ProxyLoss(temperature=args.temp)
  
      if args.cuda:
          mlp.cuda()
          embed.cuda()
          model.cuda()
          proxies.cuda()
  
      parameters_set = []
  
      layers = []
      for c in mlp.children():
          if isinstance(c, nn.Linear) or isinstance(c, nn.ModuleList):
              layers.extend(list(c.parameters()))
  
      for c in embed.children():
          if isinstance(c, nn.Linear):
              layers.extend(list(c.parameters()))
  
      parameters_set.append({'params': layers,
                             'lr': args.lr})
  
      optimizer = optim.SGD(
          parameters_set,
          lr=args.lr,
          momentum=0.9, nesterov=True,
          weight_decay=5e-5)
  
      n_parameters = sum([p.data.nelement()
                          for p in model.parameters()])
      print('  + Number of params: {}'.format(n_parameters))
  
      scheduler = CosineAnnealingLR(optimizer, args.epochs)
  
      best_acc = 0
      print('Random results:')
      if args.validation:
          validate(val_seen_loader, val_unseen_loader,
                   model, criterion)
      else:
          validate(test_seen_loader, test_unseen_loader,
                   model, criterion)
  
      for epoch in range(args.start_epoch, args.epochs + 1):
          # update learning rate
          if args.lr_decay:
              scheduler.step()
  
          # train for one epoch
          if args.validation:
              train(train_loader, model, criterion, optimizer, epoch)
              validate(val_seen_loader, val_unseen_loader,
                       model, criterion)
          else:
              train(trainval_loader, model, criterion, optimizer, epoch)
              validate(test_seen_loader, test_unseen_loader,
                       model, criterion)
  
          # saving
          save_checkpoint({
              'epoch': epoch,
              'state_dict': model.state_dict()})
  
      print('\nFinal evaluation on last epoch model:')
      validate(test_seen_loader, test_unseen_loader, model, criterion)
  
  
  def train(train_loader, model, criterion, optimizer, epoch):
      """Training loop for one epoch"""
      batch_time = AverageMeter()
      data_time = AverageMeter()
      val_losses = AverageMeter()
      val_acc = AverageMeter()
  
      # switch to train mode
      model.train()
      proxies = model.proxy_net.proxies.weight
      end = time.time()
  
      for i, (img, l) in enumerate(train_loader):
          # measure data loading time
          data_time.update(time.time() - end)
  
          if args.gan_path is None:
              labels = l
          else:
              labels, m_s, m_u = l
  
          if args.cuda:
              img = img.cuda(non_blocking=True)
              labels = labels.cuda(non_blocking=True)
              if args.gan_path is not None:
                  m_s = m_s.cuda(non_blocking=True)
                  m_u = m_u.cuda(non_blocking=True)
  
          # forward pass
          output = model(img)
  
          # compute losses
          loss, acc = criterion(output, labels, proxies)
  
          if args.gan_path is not None:
              probs = criterion.classify(output, proxies)
              seen_idx = train_loader.dataset.data['seen_Y']
              unseen_idx = train_loader.dataset.data['unseen_Y']
  
              entropy_seen = - (torch.log(probs[:, seen_idx]) * probs[:, seen_idx]).sum(1)
              entropy_unseen = - (torch.log(probs[:, unseen_idx]) * probs[:, unseen_idx]).sum(1)
  
              margin = args.margin
  
              seen_in = (entropy_seen * m_s) / (m_s.sum() + 1e-8)
              seen_out = (entropy_seen * m_u) / (m_u.sum() + 1e-8)
              seen_entropy = torch.relu(margin + seen_in.sum() - seen_out.sum())
  
              unseen_out = (entropy_unseen * m_s) / (m_s.sum() + 1e-8)
              unseen_in = (entropy_unseen * m_u) / (m_u.sum() + 1e-8)
              unseen_entropy = torch.relu(margin + unseen_in.sum() - unseen_out.sum())
  
              loss = loss + args.ent * (seen_entropy + unseen_entropy)
  
          val_losses.update(loss.item(), img.size(0))
          val_acc.update(acc, img.size(0))
  
          # 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()
  
      txt = ('\nEpoch [%02d]: Time %.3f\t Data %.3f\t'
             'Loss %.4f\t Acc %.4f' % \
             (epoch, batch_time.avg * i, data_time.avg * i,
              val_losses.avg, val_acc.avg * 100.))
      print(txt)
      write_logs(txt)
  
  
  def extract(loader, model, criterion, proxies):
  
      proxies_s, proxies_u, proxies_f = proxies
  
      probs_s = []
      probs_u = []
      probs_f = []
  
      with torch.no_grad():
          for i, (img, labels) in enumerate(loader):
              if args.cuda:
                  img = img.cuda(non_blocking=True)
  
              # compute output
              output = model(img)
  
              probs_s.append(criterion.classify(output, proxies_s).cpu())
              probs_u.append(criterion.classify(output, proxies_u).cpu())
              probs_f.append(criterion.classify(output, proxies_f).cpu())
  
      # numpy plumbery
      probs_s = np.vstack(probs_s)
      probs_u = np.vstack(probs_u)
      probs_f = np.vstack(probs_f)
  
      return [probs_s, probs_u, probs_f]
  
  
  def validate(seen_loader, unseen_loader, model, criterion):
      # switch to evaluate mode
      model.eval()
  
      # proxies
      if args.sentences:
          key = 'sentences'
      else:
          key = 'emb'
      proxies_s = seen_loader.dataset.data[key].T
      proxies_u = unseen_loader.dataset.data[key].T
  
      key = 'full_' + key
      proxies_f = seen_loader.dataset.data[key].T
  
      proxies_s = torch.from_numpy(proxies_s).cuda()
      proxies_u = torch.from_numpy(proxies_u).cuda()
      proxies_f = torch.from_numpy(proxies_f).cuda()
      proxies = [proxies_s, proxies_u, proxies_f]
  
      probs_seen = extract(seen_loader, model, criterion, proxies)
      probs_unseen = extract(unseen_loader, model, criterion, proxies)
  
      seen_idx = np.asarray([*seen_loader.dataset.data['classnames'].keys()])
  
      calibrate(probs_seen[-1], probs_unseen[-1],
                seen_loader.dataset.data['Y_orig'],
                unseen_loader.dataset.data['Y_orig'],
                seen_idx, v=write_logs)
  
  
  def save_checkpoint(state, folder=path, filename='checkpoint.pth.tar'):
      torch.save(state, os.path.join(path, filename))
  
  
  def to_numpy(x):
      return x.cpu().data.numpy()[0]
  
  
  if __name__ == '__main__':
      main()
  