text: from __future__ import print_function
  
  import argparse
  import os
  import time
  
  import model
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  import torch.optim.lr_scheduler as lr_scheduler
  from datasets.mnist import getmnist, NUM_TRAIN
  from torchvision import datasets, transforms
  
  from helper import accuracy, AverageMeter, save_checkpoint
  from module.layer1 import crxb_Conv2d, crxb_Linear
  
  
  # import pydevd_pycharm
  # pydevd_pycharm.settrace('0.0.0.0', port=12346, stdoutToServer=True, stderrToServer=True)
  import visdom
  vis = visdom.Visdom(env ="weight_statistic")
  
  class Net(nn.Module):
      def __init__(self, **crxb_cfg):
          super(Net, self).__init__()
          # self.conv1 = nn.Conv2d(1, 20, 5, 1)
          # self.conv2 = nn.Conv2d(20, 50, 5, 1)
          # self.fc1 = nn.Linear(4*4*50, 500)
          # self.fc2 = nn.Linear(500, 10)
          self.conv1 = crxb_Conv2d(1, 20, kernel_size=5, is_first_layer=True, **crxb_cfg)
          self.conv2 = crxb_Conv2d(20, 50, kernel_size=5, **crxb_cfg)
          self.conv2_drop = nn.Dropout2d()
          self.fc1 = crxb_Linear(4 * 4 * 50, 500, **crxb_cfg)
          self.fc2 = crxb_Linear(500, 10, is_last_layer=True, **crxb_cfg)
          self.activation_name = ["conv1", "max_pool2d", "conv2", "conv2_drop", "max_pool2d", "fc1", "dropout", "fc2"]
  
  
      def forward(self, x):
          self.activation = [x]
          x = F.relu(self.conv1(x))  # leaky_relu
          self.activation.append(x)
  
          x = F.max_pool2d(x, 2, 2)
          self.activation.append(x)
  
          x = F.relu(self.conv2(x))
          self.activation.append(x)
  
          x = self.conv2_drop(x)
          self.activation.append(x)
  
          x = F.max_pool2d(x, 2, 2)
          self.activation.append(x)
  
          x = x.view(-1, 4 * 4 * 50)
          x = F.relu(self.fc1(x))
          self.activation.append(x)
  
          x = F.dropout(x, training=self.training)
          self.activation.append(x)
  
          x = self.fc2(x)
          self.activation.append(x)
          return F.log_softmax(x, dim=1)
  
      def plot_histogram(self):
          for i, layer in enumerate(self.children()):
              if isinstance(layer, (crxb_Conv2d,crxb_Linear)):
                  data = layer.weight.data.view(-1)
                  vis.histogram(data, 'w%d'%i, opts=dict(title='w%d'%i))
                  data = layer.weight_quan.data.view(-1)*layer.delta_w
                  vis.histogram(data, 'wq%d'%i, opts=dict(title='wq%d'%i))
  
  # Define loss, optimizer, scheduler
  criterion = nn.CrossEntropyLoss()
  
  
  def test(args, model, device, test_loader):
      if test_loader.dataset.train:
          print("test on validation set\r\n")
      else:
          print("test on test set\r\n")
  
      # validate
      model.eval()
      # test_loss = 0
      # correct = 0
      # num_samples = 0
      losses = AverageMeter()
      top1 = AverageMeter()
      top5 = AverageMeter()
      with torch.no_grad():
          for data in test_loader:
              inputs, labels = data[0].to(device), data[1].to(device)
              outputs = model(inputs)
              test_loss = criterion(outputs, labels).item()
              # pred = outputs.argmax(dim=1, keepdim=True)
              # correct += pred.eq(labels.view_as(pred)).sum().item()
              # # _, pred = torch.max(outputs, 1)
              # # correct += (pred == labels).sum().item()
              # num_samples += pred.size(0)
              prec1, prec5 = accuracy(outputs, labels, topk=(1, 5))
              losses.update(test_loss, labels.size(0))
              top1.update(prec1[0], labels.size(0))
              top5.update(prec5[0], labels.size(0))
  
      print('\nTest set: Average loss: {:.4f}, Accuracy: Prec@1:{}/{} ({:.2f}%) Prec@5:{}/{} ({:.2f}%)\n'.format(
          losses.avg, top1.sum // 100, top1.count, top1.avg, top5.sum // 100, top1.count, top5.avg))
      with torch.no_grad():
          # model.clip_w()
          model.plot_histogram()
          # model.plot_qcurve()
      return top1.avg, top5.avg, losses.avg
  
  
  best_prec1 = 0
  
  
  def main():
      # Training settings
      parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
      # model cfg
      parser.add_argument('--model-type', type=str, default="MNIST", help="type of the model.")
      parser.add_argument('--model-structure', type=int, default=0, metavar='N',
                          help='model structure to be trained (default: 0)')
      parser.add_argument('--resume', default='', type=str, metavar='PATH',
                          help='path to latest checkpoint, (default: None)')
      parser.add_argument('--e', '--evaluate', dest='evaluate', action='store_true',
                          help='evaluate model on validation set')
  
      # dataset
      parser.add_argument('--dataset-root', type=str, default="../datasets", help="load dataset path.")
      parser.add_argument('--workers', default=0, type=int, metavar='N',
                          help='number of data loading workers (default: 0)')
      parser.add_argument('--train-batch-size', type=int, default=64, metavar='N',
                          help='input batch size for training (default: 64)')
      parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                          help='input batch size for testing (default: 1000)')
      # train cfg
      parser.add_argument('--epochs', type=int, default=20, metavar='N',
                          help='number of epochs to train (default: 10)')
      parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                          help='manual epoch number (useful to restarts)')
      parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                          help='learning rate (default: 0.01)')
  
      # device init cfg
      parser.add_argument('--no-cuda', action='store_true', default=False,
                          help='disables CUDA training')
      parser.add_argument('--seed', type=int, default=1, metavar='S',
                          help='random seed (default: 1)')
  
      # optimizer
      parser.add_argument('--optim', type=str, default="", help="optim type Adam/SGD")
      parser.add_argument('--resume-optim', action='store_true', default=False,
                          help='resume optim')
      parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                          help='SGD momentum (default: 0.5)')
      parser.add_argument('--wd', default=5e-4, type=float,
                          metavar='W', help='weight decay (default: 5e-4)')
      # scheduler
      parser.add_argument('--scheduler', type=str, default="None", help="scheduler MultiStepLR/None/ReduceLROnPlateau")
      parser.add_argument('--gamma', type=float, default=0.1, help='LR is multiplied by gamma on schedule.')
      parser.add_argument('--decreasing-lr', default='10', help='decreasing strategy')
  
      # result output cfg
      parser.add_argument('--detail', action='store_true', default=False,
                          help='show log in detial')
      parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                          help='how many batches to wait before logging training status')
      parser.add_argument('--save-model', action='store_true', default=True,
                          help='For Saving the current Model')
      parser.add_argument('--checkpoint-path', type=str, default="", help="save model path.")
  
      # crossbar cfg
      parser.add_argument('--Quantized', action='store_true', default=False,
                          help='use quantized model')
      parser.add_argument('--qbit', default='8', help='activation/weight qbit')
  
      parser.add_argument('--crxb-size', type=int, default=64, help='corssbar size')
      parser.add_argument('--vdd', type=float, default=3.3, help='supply voltage')
      parser.add_argument('--gwire', type=float, default=0.0357,
                          help='wire conductacne')
      parser.add_argument('--gload', type=float, default=0.25,
                          help='load conductance')
      parser.add_argument('--gmax', type=float, default=0.000333,
                          help='maximum cell conductance')
      parser.add_argument('--gmin', type=float, default=0.000000333,
                          help='minimum cell conductance')
      parser.add_argument('--ir-drop', action='store_true', default=False,
                          help='switch to turn on ir drop analysis')
      parser.add_argument('--scaler-dw', type=float, default=1,
                          help='scaler to compress the conductance')
      parser.add_argument('--test', action='store_true', default=False,
                          help='switch to turn inference mode')
      parser.add_argument('--enable_noise', action='store_true', default=False,
                          help='switch to turn on noise analysis')
      parser.add_argument('--enable_SAF', action='store_true', default=False,
                          help='switch to turn on SAF analysis')
      parser.add_argument('--enable_ec-SAF', action='store_true', default=False,
                          help='switch to turn on SAF error correction')
      parser.add_argument('--freq', type=float, default=10e6,
                          help='scaler to compress the conductance')
      parser.add_argument('--temp', type=float, default=300,
                          help='scaler to compress the conductance')
  
      args = parser.parse_args()
      print("+++", args)
  
      # Train the network on the training data
      # Test the network on the test data
      use_cuda = not args.no_cuda and torch.cuda.is_available()
      torch.manual_seed(args.seed)
      device = torch.device("cuda" if use_cuda else "cpu")
      print(device)
  
      qbit_list = list(map(int, args.qbit.split(',')))
      crxb_cfg = {'ir_drop': args.ir_drop, 'device': device,
                  'gmax': args.gmax, 'gmin': args.gmin, 'gwire': args.gwire, 'gload': args.gload,
                  'input_qbit': qbit_list[0], 'weight_qbit': qbit_list[1], 'activation_qbit': qbit_list[2],
                  'vdd': args.vdd, 'enable_noise': args.enable_noise,
                  'freq': args.freq, 'temp': args.temp, 'crxb_size': args.crxb_size,
                  'enable_SAF': args.enable_SAF, 'enable_ec_SAF': args.enable_ec_SAF}
  
      net = Net(**crxb_cfg).to(device)
  
      if torch.cuda.device_count() > 1:
          print("Let's use", torch.cuda.device_count(), "GPUs!")
          net = nn.DataParallel(net)
      net.to(device)
  
      # for param in net.parameters():
      #     param = nn.init.normal_(param)
  
      # config
      milestones = list(map(int, args.decreasing_lr.split(',')))
      print(milestones)
  
      # optimizer = optim.SGD(net.parameters(), lr=lr, momentum=MOMENTUM, weight_decay=WEIGHT_DECAY) # not good enough 68%
      # optimizer = optim.Adam(net.parameters(), lr=args.lr, weight_decay=args.wd)
      if args.optim == "Adam":
          optimizer = optim.Adam(net.parameters(), lr=args.lr, weight_decay=args.wd)
      elif args.optim == "SGD":
          optimizer = optim.SGD(net.parameters(), lr=args.lr, weight_decay=args.wd, momentum=args.momentum)
      else:
          optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum)
  
      # optionlly resume from a checkpoint
      if args.resume:
          print("=> using pre-trained model '{}'".format(args.model_type))
      else:
          print("=> creating model '{}'".format(args.model_type))
  
      global best_prec1
      if args.resume:
          if os.path.isfile(args.resume):
              print("=> loading checkpoint '{}'".format(args.resume))
  
              checkpoint = torch.load(args.resume)
              args.start_epoch = checkpoint['epoch']
              best_prec1 = checkpoint['best_prec1']
  
              model_dict = net.state_dict()
              pretrained_dict = torch.load(args.resume)['state_dict']
              # pretrained_dict = torch.load(args.resume,map_location=torch.device('cpu'))['state_dict']
              print(model_dict.keys(),'\r\n')
              print(pretrained_dict.keys(),'\r\n')
              new_dict = {}
              for k,v in model_dict.items():
                  if 'module.'+k in pretrained_dict:
                      new_dict[k] = pretrained_dict['module.'+k]
                      print(k,' ')
                  else:
                      new_dict[k] = pretrained_dict[k]
                      print(k,'!!')
              model_dict.update(new_dict)
              net.load_state_dict(model_dict)
  
              # net.load_state_dict(checkpoint['state_dict'])
              if args.resume_optim:
                  try:
                      optimizer.load_state_dict(checkpoint['optimizer'])
                  except KeyError:
                      print("saved optim not compatible")
              print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
          else:
              print("=> no checkpoint found at '{}'".format(args.resume))
  
  
      # Data loading
      kwargs = {'num_workers': args.workers, 'pin_memory': True} if use_cuda else {}
      trainloader = torch.utils.data.DataLoader(
          datasets.MNIST(args.dataset_root, train=True, download=True,
                         transform=transforms.Compose([
                             transforms.ToTensor(),
                             transforms.Normalize((0.1307,), (0.3081,))
                         ])),
          batch_size=args.train_batch_size, shuffle=True, **kwargs)
  
      testloader = torch.utils.data.DataLoader(
          datasets.MNIST(args.dataset_root, train=False, download=True, transform=transforms.Compose([
              transforms.ToTensor(),
              transforms.Normalize((0.1307,), (0.3081,))
          ])),
          batch_size=args.test_batch_size, shuffle=True, **kwargs)
      print(len(trainloader), len(testloader))
  
      t_s = time.monotonic()
  
      print("!!test!!")
      test(args, net, device, testloader)
      t_e = time.monotonic()
  
      m, s = divmod(t_e - t_s, 60)
      h, m = divmod(m, 60)
      print("%d:%02d:%02d" % (h, m, s))
  
      # print([i for i in net.named_modules()],len([i for i in net.named_modules()]),"\n")
      # print([i for i in net.named_children()],len([i for i in net.named_children()]),"\n")
      # plot_histogram(net.activation, net.activation_name)
      # for epoch in range(1, args.epochs + 1):
      #     test(args, net, device, test_loader)
      #     plot_histogram(net.activation)
  
      parm_name = []
      parm_value = []
      for name, parameters in net.named_parameters():
          print(name, ':', parameters.size())
          parm_name.append(name)
          parm_value.append(parameters.detach())
  
      # plot_histogram(parm_value, parm_name)
  
      from scipy import stats
      for weight in parm_value:
          N = weight.cpu().numpy().ravel().shape[0]
          p = 0.03
          acc_p = 0
          print('length:%d\n'%N)
          k = stats.binom.ppf(p,N,p)
          # stats.truncnorm
          print(N-k,stats.binom.cdf(k,N,p))
          # while acc_p<0.03:
          #     acc_p =+ comb(N, k)*(p**k)*((1-p)**(N-k))
          #     if np.isnan(acc_p):
          #         break
          #     k += 1
          # if np.isnan(acc_p):
          #     k = stats.poisson.ppf(p, N*p)
          #     print(k, stats.poisson.cdf(k, N*p))
          #     break
          # else:
          #     print(k, acc_p,'\n')
  
  
  
  
  if __name__ == '__main__':
      main()
  
  # 2020/03/25: divide the quatization bits into three class, IA, W, MA.
  # format the physical parameter dict: crxb_cfg
  # add layer number flag: is_first_layer, is_last_layer
  # 2020/03/25: add layer number flag: is_first_layer, is_last_layer
  