text: from __future__ import print_function
  
  import argparse
  import itertools
  import os
  import pickle
  import sys
  from datetime import datetime
  
  import matplotlib
  import numpy as np
  import torch
  
  matplotlib.use('Agg')
  import matplotlib.pyplot as plt
  from proj import archs
  from proj.utils.cluster.general import config_to_str, get_opt, \
      update_lr, nice
  from proj.utils.cluster.cluster_eval import cluster_eval
  from proj.utils.cluster.data import cluster_create_dataloaders
  from proj.utils.cluster.IID_losses import IID_loss
  
  """
    Semisupervised overclustering ("IIC+" = "IID+")
    Note network is trained entirely unsupervised, as labels are found for 
    evaluation only and do not affect the network.
    Train and test script (greyscale datasets).
    Network has one output head only.
  """
  
  # Options ----------------------------------------------------------------------
  
  parser = argparse.ArgumentParser()
  parser.add_argument("--model_ind", type=int, required=True)
  parser.add_argument("--arch", type=str, default="ClusterNet4h")
  parser.add_argument("--opt", type=str, default="Adam")
  parser.add_argument("--mode", type=str, default="IID+")
  
  parser.add_argument("--dataset", type=str, default="MNIST")
  parser.add_argument("--dataset_root", type=str,
                      default="MNIST")
  
  parser.add_argument("--gt_k", type=int, default=10)
  parser.add_argument("--output_k", type=int, default=10)
  parser.add_argument("--lamb", type=float, default=1.0)
  parser.add_argument("--lr", type=float, default=0.01)
  parser.add_argument("--lr_schedule", type=int, nargs="+", default=[])
  parser.add_argument("--lr_mult", type=float, default=0.1)
  
  parser.add_argument("--num_epochs", type=int, default=1000)
  parser.add_argument("--batch_sz", type=int, default=240)  # num pairs
  parser.add_argument("--num_dataloaders", type=int, default=3)
  parser.add_argument("--num_sub_heads", type=int, default=5)
  
  parser.add_argument("--out_root", type=str,
                      default="iid_private")
  parser.add_argument("--restart", dest="restart", default=False,
                      action="store_true")
  parser.add_argument("--restart_from_best", dest="restart_from_best",
                      default=False, action="store_true")
  parser.add_argument("--test_code", dest="test_code", default=False,
                      action="store_true")
  
  parser.add_argument("--save_freq", type=int, default=10)
  
  parser.add_argument("--batchnorm_track", default=False, action="store_true")
  
  # transforms
  parser.add_argument("--mix_train", dest="mix_train", default=False,
                      action="store_true")
  parser.add_argument("--demean", dest="demean", default=False,
                      action="store_true")
  parser.add_argument("--per_img_demean", dest="per_img_demean", default=False,
                      action="store_true")
  parser.add_argument("--data_mean", type=float, nargs="+",
                      default=[0.5, 0.5, 0.5])
  parser.add_argument("--data_std", type=float, nargs="+",
                      default=[0.5, 0.5, 0.5])
  
  parser.add_argument("--crop_orig", dest="crop_orig", default=False,
                      action="store_true")
  parser.add_argument("--crop_other", dest="crop_other", default=False,
                      action="store_true")
  parser.add_argument("--tf1_crop", type=str, default="random")  # type name
  parser.add_argument("--tf2_crop", type=str, default="random")
  parser.add_argument("--tf1_crop_sz", type=int, default=84)
  parser.add_argument("--tf2_crop_szs", type=int, nargs="+",
                      default=[84])  # allow diff crop for imgs_tf
  parser.add_argument("--tf3_crop_diff", dest="tf3_crop_diff", default=False,
                      action="store_true")
  parser.add_argument("--tf3_crop_sz", type=int, default=0)
  parser.add_argument("--input_sz", type=int, default=96)
  
  parser.add_argument("--rot_val", type=float, default=0.)
  parser.add_argument("--always_rot", dest="always_rot", default=False,
                      action="store_true")
  parser.add_argument("--no_jitter", dest="no_jitter", default=False,
                      action="store_true")
  parser.add_argument("--no_flip", dest="no_flip", default=False,
                      action="store_true")
  
  config = parser.parse_args()
  
  # Setup ------------------------------------------------------------------------
  
  config.twohead = False
  
  config.in_channels = 1
  config.out_dir = os.path.join(config.out_root, str(config.model_ind))
  assert (config.batch_sz % config.num_dataloaders == 0)
  config.dataloader_batch_sz = int(config.batch_sz / config.num_dataloaders)
  
  assert (config.mode == "IID+")
  assert (config.output_k >= config.gt_k)
  config.eval_mode = "orig"
  config.double_eval = False
  
  if "MNIST" == config.dataset:
      config.train_partitions = [True]
      config.mapping_assignment_partitions = [True]
      config.mapping_test_partitions = [False]
  
  if not os.path.exists(config.out_dir):
      os.makedirs(config.out_dir)
  
  if config.restart:
      config_name = "config.pickle"
      net_name = "latest_net.pytorch"
      opt_name = "latest_optimiser.pytorch"
  
      if config.restart_from_best:
          config_name = "best_config.pickle"
          net_name = "best_net.pytorch"
          opt_name = "best_optimiser.pytorch"
  
      given_config = config
      reloaded_config_path = os.path.join(given_config.out_dir, config_name)
      print("Loading restarting config from: %s" % reloaded_config_path)
      with open(reloaded_config_path, "rb") as config_f:
          config = pickle.load(config_f)
      assert (config.model_ind == given_config.model_ind)
      config.restart = True
      config.restart_from_best = given_config.restart_from_best
  
      config.num_epochs = given_config.num_epochs
      config.lr_schedule = given_config.lr_schedule
  
  else:
      print("Config: %s" % config_to_str(config))
  
  # Model ------------------------------------------------------------------------
  
  dataloaders, mapping_assignment_dataloader, mapping_test_dataloader = \
      cluster_create_dataloaders(config)
  
  net = archs.__dict__[config.arch](config)
  if config.restart:
      model_path = os.path.join(config.out_dir, net_name)
      net.load_state_dict(
          torch.load(model_path, map_location=lambda storage, loc: storage))
  
  net.cuda()
  net = torch.nn.DataParallel(net)
  net.train()
  
  optimiser = get_opt(config.opt)(net.module.parameters(), lr=config.lr)
  if config.restart:
      print("loading latest opt")
      optimiser.load_state_dict(
          torch.load(os.path.join(config.out_dir, opt_name)))
  
  if config.restart:
      if not config.restart_from_best:
          next_epoch = config.last_epoch + 1  # corresponds to last saved model
      else:
          next_epoch = np.argmax(np.array(config.epoch_acc)) + 1
      print("starting from epoch %d" % next_epoch)
  
      config.epoch_acc = config.epoch_acc[:next_epoch]  # in case we overshot
      config.epoch_avg_subhead_acc = config.epoch_avg_subhead_acc[:next_epoch]
      config.epoch_stats = config.epoch_stats[:next_epoch]
  
      config.epoch_loss = config.epoch_loss[:(next_epoch - 1)]
      config.epoch_loss_no_lamb = config.epoch_loss_no_lamb[:(next_epoch - 1)]
  else:
      config.epoch_acc = []
      config.epoch_avg_subhead_acc = []
      config.epoch_stats = []
  
      config.epoch_loss = []
      config.epoch_loss_no_lamb = []
  
      _ = cluster_eval(config, net,
                       mapping_assignment_dataloader=mapping_assignment_dataloader,
                       mapping_test_dataloader=mapping_test_dataloader,
                       sobel=False)
  
      print("Pre: time %s: \n %s" % (datetime.now(), nice(config.epoch_stats[-1])))
      sys.stdout.flush()
      next_epoch = 1
  
  fig, axarr = plt.subplots(4, sharex=False, figsize=(20, 20))
  
  # Train ------------------------------------------------------------------------
  
  for e_i in range(next_epoch, config.num_epochs):
      print("Starting e_i: %d" % e_i)
      sys.stdout.flush()
  
      iterators = (d for d in dataloaders)
  
      b_i = 0
  
      if e_i in config.lr_schedule:
          optimiser = update_lr(optimiser, lr_mult=config.lr_mult)
  
      avg_loss = 0.  # over epoch
      avg_loss_no_lamb = 0.
      avg_loss_count = 0
  
      for tup in itertools.izip(*iterators):
          net.module.zero_grad()
  
          all_imgs = torch.zeros(config.batch_sz, config.in_channels,
                                 config.input_sz,
                                 config.input_sz).cuda()
          all_imgs_tf = torch.zeros(config.batch_sz, config.in_channels,
                                    config.input_sz,
                                    config.input_sz).cuda()
  
          imgs_curr = tup[0][0]  # always the first
          curr_batch_sz = imgs_curr.size(0)
          for d_i in range(config.num_dataloaders):
              imgs_tf_curr = tup[1 + d_i][0]  # from 2nd to last
              assert (curr_batch_sz == imgs_tf_curr.size(0))
  
              actual_batch_start = d_i * curr_batch_sz
              actual_batch_end = actual_batch_start + curr_batch_sz
              all_imgs[actual_batch_start:actual_batch_end, :, :, :] = \
                  imgs_curr.cuda()
              all_imgs_tf[actual_batch_start:actual_batch_end, :, :, :] = \
                  imgs_tf_curr.cuda()
  
          if not (curr_batch_sz == config.dataloader_batch_sz):
              print("last batch sz %d" % curr_batch_sz)
  
          curr_total_batch_sz = curr_batch_sz * config.num_dataloaders  # times 2
          all_imgs = all_imgs[:curr_total_batch_sz, :, :, :]
          all_imgs_tf = all_imgs_tf[:curr_total_batch_sz, :, :, :]
  
          x_outs = net(all_imgs)
          x_tf_outs = net(all_imgs_tf)
  
          avg_loss_batch = None  # avg over the heads
          avg_loss_no_lamb_batch = None
          for i in range(config.num_sub_heads):
              loss, loss_no_lamb = IID_loss(x_outs[i], x_tf_outs[i], lamb=config.lamb)
              if avg_loss_batch is None:
                  avg_loss_batch = loss
                  avg_loss_no_lamb_batch = loss_no_lamb
              else:
                  avg_loss_batch += loss
                  avg_loss_no_lamb_batch += loss_no_lamb
  
          avg_loss_batch /= config.num_sub_heads
          avg_loss_no_lamb_batch /= config.num_sub_heads
  
          if ((b_i % 100) == 0) or (e_i == next_epoch):
              print("Model ind %d epoch %d batch: %d avg loss %f avg loss no lamb %f "
                    "time %s" % \
                    (config.model_ind, e_i, b_i, avg_loss_batch.item(),
                     avg_loss_no_lamb_batch.item(), datetime.now()))
              sys.stdout.flush()
  
          if not np.isfinite(avg_loss_batch.item()):
              print("Loss is not finite... %s:" % str(avg_loss_batch))
              exit(1)
  
          avg_loss += avg_loss_batch.item()
          avg_loss_no_lamb += avg_loss_no_lamb_batch.item()
          avg_loss_count += 1
  
          avg_loss_batch.backward()
  
          optimiser.step()
  
          b_i += 1
          if b_i == 2 and config.test_code:
              break
  
      # Eval -----------------------------------------------------------------------
  
      avg_loss = float(avg_loss / avg_loss_count)
      avg_loss_no_lamb = float(avg_loss_no_lamb / avg_loss_count)
  
      config.epoch_loss.append(avg_loss)
      config.epoch_loss_no_lamb.append(avg_loss_no_lamb)
  
      is_best = cluster_eval(config, net,
                             mapping_assignment_dataloader=mapping_assignment_dataloader,
                             mapping_test_dataloader=mapping_test_dataloader,
                             sobel=False)
  
      print("Pre: time %s: \n %s" % (datetime.now(), nice(config.epoch_stats[-1])))
      sys.stdout.flush()
  
      axarr[0].clear()
      axarr[0].plot(config.epoch_acc)
      axarr[0].set_title("acc (best), top: %f" % max(config.epoch_acc))
  
      axarr[1].clear()
      axarr[1].plot(config.epoch_avg_subhead_acc)
      axarr[1].set_title("acc (avg), top: %f" % max(config.epoch_avg_subhead_acc))
  
      axarr[2].clear()
      axarr[2].plot(config.epoch_loss)
      axarr[2].set_title("Loss")
  
      axarr[3].clear()
      axarr[3].plot(config.epoch_loss_no_lamb)
      axarr[3].set_title("Loss no lamb")
  
      fig.canvas.draw_idle()
      fig.savefig(os.path.join(config.out_dir, "plots.png"))
  
      if is_best or (e_i % config.save_freq == 0):
          net.module.cpu()
  
          if e_i % config.save_freq == 0:
              torch.save(net.module.state_dict(),
                         os.path.join(config.out_dir, "latest_net.pytorch"))
              torch.save(optimiser.state_dict(),
                         os.path.join(config.out_dir, "latest_optimiser.pytorch"))
              config.last_epoch = e_i  # for last saved version
  
          if is_best:
              torch.save(net.module.state_dict(),
                         os.path.join(config.out_dir, "best_net.pytorch"))
              torch.save(optimiser.state_dict(),
                         os.path.join(config.out_dir, "best_optimiser.pytorch"))
  
              with open(os.path.join(config.out_dir, "best_config.pickle"),
                        'wb') as outfile:
                  pickle.dump(config, outfile)
  
              with open(os.path.join(config.out_dir, "best_config.txt"),
                        "w") as text_file:
                  text_file.write("%s" % config)
  
          net.module.cuda()
  
      with open(os.path.join(config.out_dir, "config.pickle"), 'wb') as outfile:
          pickle.dump(config, outfile)
  
      with open(os.path.join(config.out_dir, "config.txt"), "w") as text_file:
          text_file.write("%s" % config)
  
      if config.test_code:
          exit(0)
  