drrunsp.py

#!/usr/bin/python3
#!/usr/bin/python2.7
#python3 drrunsp.py --gpu 1 --epochs 1000 --pt 10_100 --opt adam --lr 0.0000002 --batch_size 32 --save test35 --pool 2 --channel 0,1 --dform image --stride label --network 3 --normalize 1 --cand gamma,pi0
from __future__ import print_function
import logging
import argparse
import os,sys,json
import tensorflow.keras as keras
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import *
from numpy.random import seed
import subprocess, warnings
import random, math
import shutil, datetime
from array import array
import numpy as np
#import ROOT as rt
import tensorflow as tf
#tf.autograph.set_verbosity(5)
#tf.get_logger().setLevel('INFO')
#from keras.backend.tensorflow_backend import set_session
from importlib import import_module
from sklearn.utils import shuffle
from sklearn.preprocessing import normalize
from sklearn.metrics import roc_auc_score, auc, roc_curve
from driter import *
import matplotlib.pyplot as plt

class drrun:
  def __init__(self,args):
    seed(args.seed)
    self.start=datetime.datetime.now()
    self.args=args
    os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"]=str(args.gpu)
    if(args.gpu!=-1):
      gpus = tf.config.experimental.list_physical_devices('GPU')
      print(gpus)
      tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
      tf.config.experimental.set_memory_growth(gpus[0], True)

    self.savename='save/'+str(args.save)
    self.num_classes = 2 
    self.classes=[]
    self.data_form=args.dform
    self.candidate=args.cand.strip(" ").split(",")
    self.num_classes=len(self.candidate)
    self.sigdata=[]
    self.sigfile=[]
    self.sigshape=[]
    self.model=0
    self.classes=['el','pi','ga','pi0','ka','ka0','proton','neutron']
    self.trainY=[]
    self.trainp=[]
    self.jump=[]
    for i in self.candidate:
      npz_name="npzs/{}_{}GeV_sparse.npz".format(i,args.pt)
      #npz_name="/hdfs/ml/dualreadout/{}/{}_{}GeV_sparse.npz".format(args.version,i,args.pt)
      self.sigfile.append(npz_name)
      self.sigdata.append(np.load(npz_name,allow_pickle=True))
      self.jump.append(0)
      #self.sigdata.append(np.load(os.path.abspath(os.getcwd())+"/npzs/{}/dr{}_{}_sparse.npz".format(args.version,args.pt,i),allow_pickle=True))
    self.siglabel=[]
    for i in range(self.num_classes):
      label=[0.]*self.num_classes
      label[i]=1.
      self.siglabel.append(label)
      self.sigshape.append(self.sigdata[i]["info"].item()["{}_shape".format(self.data_form)])

    self.stride=args.stride.split(",")
    self.channel=[int(i) for i in args.channel.split(",")]
    if(args.extra==""):
      self.extra=[]
    else:
      self.extra=[int(i) for i in args.extra.split(",")]
    self.balance=min([len(self.sigdata[i]['dr_ecorr'])-self.jump[i] for i in range(self.num_classes)])
    print("balance",[len(self.sigdata[i]['dr_ecorr'])-self.jump[i] for i in range(self.num_classes)])
    self.len_train=int(self.balance*0.7)

    print(self.candidate,self.siglabel," strid ",self.stride)
    print(sys.argv)
    print("total entries ",{self.candidate[i]:len(self.sigdata[i]['dr_ecorr'])-self.jump[i] for i in range(self.num_classes)})
    print("balanced {} * num_category is number of used data".format(self.balance))
    print(self.sigfile)
    print(self)

  def premodel(self,savename=None,stride=None,num_classes=None,channel=None,extra=None,args=None):
    if(savename==None):savename=self.savename
    if(stride==None):stride=self.stride
    if(num_classes==None):num_classes=self.num_classes
    if(channel==None):channel=self.channel
    if(extra==None):extra=self.extra
    if(args==None):args=self.args
    net=import_module('symbols.symbols')
    if(args.opt=="sgd"):
      opt=keras.optimizers.SGD()
    if(args.opt=="rms"):
      opt=keras.optimizers.RMSprop()
    if(args.opt=="adam"):
      opt=keras.optimizers.Adam(lr=args.lr)
    losses={}
    metrics={}
    lossweight={}
    if(sys.version_info[0]>=3):
      model_metrics=['accuracy',keras.metrics.AUC()]
    else:
      model_metrics=['accuracy']
    for i in stride:
      if(i=="label"):
        losses["output{}".format(i)]=args.loss
        metrics["output{}".format(i)]=model_metrics
      else:
        losses["output{}".format(i)]="mse"
        metrics["output{}".format(i)]=["mse"]
      lossweight["output{}".format(i)]=1.0

    if(args.dform=="image"):
      if(args.gpu!=-1):
        tf.keras.backend.set_image_data_format('channels_first')
      else:
        tf.keras.backend.set_image_data_format('channels_last')
      if("center" in args.memo):
        self.model=net.emmodel2((len(channel),args.pix,args.pix),num_classes=num_classes)
      elif(args.pool>42):
        self.model=net.emmodel1((len(channel),args.pix,args.pix),args.pool,num_classes=num_classes,network=args.network,extra=extra)
      else:
        self.model=net.emmodel0((len(channel),args.pix,args.pix),args.pool,num_classes=num_classes,stride=stride,network=args.network,extra=extra)
    if(args.dform=="point"):
      tf.keras.backend.set_image_data_format('channels_last')
      self.model=pointmodel(args.num_point,len(channel),num_classes=num_classes,stride=stride,network=args.network)
    self.model.compile(loss=losses,
                  optimizer=opt, loss_weights=lossweight,
                  metrics=metrics)

    shutil.copy("symbols/symbols.py",savename+'/')
    shutil.copy(__file__,savename+'/')
    #from keras.utils import plot_model
    #plot_model(model,to_file=savename+'/model.png')
    #plot_model(model,to_file='/home/yulee/keras/model.png')
    print(self.model.summary())

  def getdata(self,data_form=None,num_classes=None,channel=None,extra=None,begin=0,end=-1,args=None,energy=[-1,1000000]):
    if(data_form==None):data_form=self.data_form
    if(num_classes==None):num_classes=self.num_classes
    if(channel==None):channel=self.channel
    if(extra==None):extra=self.extra
    if(args==None):args=self.args
    len_train=end-begin
    energy_pass=[]
    tot=[self.sigdata[i]['dr_ecorr'][begin+self.jump[i]:end] for i in range(num_classes)]
    if(3 in extra or 4 in extra or energy[0]!=-1):
      if(energy[0]!=-1):
        for j in range(len(tot)):
          energy_pass.append([])
          for i in range(len(tot[j])):
            if(tot[j][i]>=energy[0] and tot[j][i]<energy[1]):
              energy_pass[j].append(i+begin+self.jump[j])
        tot=[self.sigdata[i]['dr_ecorr'][energy_pass[i]] for i in range(len(energy_pass))]
    tot=np.concatenate(tot)
        #tot=np.concatenate([self.sigdata[i]['dr_ecorr'][begin+self.jump[i]:end] for i in range(num_classes)]).reshape(-1,1)

    if(data_form=="point"):
        if(energy[0]!=-1):
          buf_shape=[(len(energy_pass[i]),self.sigshape[i][1],self.sigshape[i][2]) for i in range(num_classes)]
          X=np.concatenate([self.sigdata[i][data_form].item()[energy_pass[i]].toarray().reshape(buf_shape[i])[:,:args.num_point,channel] for i in range(num_classes)])
        else:
          buf_shape=[(end-begin-self.jump[i],self.sigshape[i][1],self.sigshape[i][2]) for i in range(num_classes)]
          X=np.concatenate([self.sigdata[i][data_form].item()[begin+self.jump[i]:end].toarray().reshape(buf_shape[i])[:,:args.num_point,channel] for i in range(num_classes)])
    elif(data_form=="image"):
        if(energy[0]!=-1):
          buf_shape=[(len(energy_pass[i]),self.sigshape[i][1],self.sigshape[i][2],self.sigshape[i][3]) for i in range(num_classes)]
          X=np.concatenate([self.sigdata[i][data_form].item()[energy_pass[i]].toarray().reshape(buf_shape[i])[:,channel] for i in range(num_classes)])
        else:
          buf_shape=[(end-begin-self.jump[i],self.sigshape[i][1],self.sigshape[i][2],self.sigshape[i][3]) for i in range(num_classes)]
          X=np.concatenate([self.sigdata[i][data_form].item()[begin+self.jump[i]:end].toarray().reshape(buf_shape[i])[:,channel] for i in range(num_classes)])
    if(args.stride=='E_Gen'):
        if(energy[0]!=-1):
          Y=np.concatenate([self.sigdata[i]['E_Gen'][energy_pass[i]] for i in range(num_classes)])
        else:
          Y=np.concatenate([self.sigdata[i]['E_Gen'][begin+self.jump[i]:end] for i in range(num_classes)])
    else:
        if(energy[0]!=-1):
          Y=np.concatenate([[self.siglabel[i]]*(len(energy_pass[i])) for i in range(num_classes)])
        else:
          Y=np.concatenate([[self.siglabel[i]]*(len_train-self.jump[i]) for i in range(num_classes)])
    if(extra!=[]):
      #dr_ecorr,E_Gen,width_Gen,SC_ratio,CS_ratio,info=info)
      ex=[]
      vlist=['dr_ecorr','SC_ratio','CS_ratio']
      for j in range(len(vlist)):
        if(j in extra):
          if(energy[0]!=-1):
            ex.append(np.concatenate([self.sigdata[i][vlist[j]][energy_pass[i]] for i in range(num_classes)]).reshape(-1,1))
          else:
            ex.append(np.concatenate([self.sigdata[i][vlist[j]][begin+self.jump[i]:end] for i in range(num_classes)]).reshape(-1,1))
      if(3 in extra):
        hot=[]
        for i in range(len(tot)):
          hot.append([np.max(X[i][0])/tot[i]])
        ex.append(np.array(hot))
      if(4 in extra):
        hot=[]
        for i in range(len(tot)):
          hot.append([np.max(X[i][1])/tot[i]])
        ex.append(np.array(hot))
      for i in ex:
        print(i.shape)
      ex=np.concatenate(ex,axis=1)
      print("extra shape ",ex.shape)
      X,Y,ex,tot= shuffle(X,Y,ex,tot)
    else:
      X,Y,tot= shuffle(X,Y,tot)
    #print("##",len(X))
    #print("##",np.array(X).shape)
    print("input shape ",X.shape,"channel 0",X[:,0].shape," y shape ",Y.shape, "range {} ~ {}".format(begin,end))
    if(args.normalize==1):
      xshape=X.shape
      print('normalize',xshape)
      if(data_form=="point"):
        X=np.stack([normalize(X[:,:,i],axis=1,norm='l1').reshape(xshape[0],xshape[1]) if i > 2 else X[:,:,i] for i in channel],axis=-1)
      elif(data_form=="image"):
        X=np.stack([normalize(X[:,i].reshape(xshape[0],-1),axis=1,norm='l1').reshape(xshape[0],xshape[2],xshape[3]) for i in channel],axis=1)
    if(extra!=[]):
      X=[X,ex]
    return X,Y,tot

  def train(self,savename=None,data_form=None,num_classes=None,candidate=None,channel=None,extra=None,balance=None,len_train=None,args=None,energy=[-1,1000000]):
    if(savename==None):savename=self.savename
    if(data_form==None):data_form=self.data_form
    if(num_classes==None):num_classes=self.num_classes
    if(candidate==None):candidate=self.candidate
    if(channel==None):channel=self.channel
    if(extra==None):extra=self.extra
    if(balance==None):balance=self.balance
    if(len_train==None):len_train=self.len_train
    if(args==None):args=self.args
    X,self.trainY,self.trainE=self.getdata(data_form,num_classes,channel,extra,0,len_train,args,energy)
    if(args.job=='train'):
      checkpoint=keras.callbacks.ModelCheckpoint(filepath=savename+'/check_{epoch}',monitor='val_loss',verbose=0,save_weights_only=False,save_best_only=False,mode='auto',save_freq="epoch")
      print(self.trainY[:10])
      history=self.model.fit(X,self.trainY,batch_size=args.batch_size,epochs=args.epochs,validation_split=0.3,verbose=1,callbacks=[checkpoint])

      f=open(savename+'/history','w')
      one=history.history['val_loss'].index(min(history.history['val_loss']))
      f.write(str(one)+'\n')
      for i in range(args.epochs):
        try:
          if(i!=one):shutil.rmtree(savename+"/check_"+str(i+1),ignore_errors=True)
        except:
          print("failed to drop")
      f.write(str(history.history))
      f.close()

      print (datetime.datetime.now()-self.start)
      logging.info("memo "+args.memo)
      logging.info("spent time "+str(datetime.datetime.now()-self.start))
      logging.info("python {}".format(str(args)))
      logging.info("balance {} len_train {}".format(balance,len_train))
    self.loadmodel()
    self.trainp=self.model.predict(X,verbose=0)
    del X

  def loadmodel(self,one=-1):
    if(one==-1):
      f=open(self.savename+'/history','r')
      one=int(f.readline().strip("\n"))+1
      f.close()
    self.model=keras.models.load_model(self.savename+"/check_"+str(one))
    print("epoch {} loaded".format(one))

  def test(self,savename=None,data_form=None,num_classes=None,candidate=None,balance=None,len_train=None,channel=None,extra=None,args=None,energy=[-1,1000000]):
    if(savename==None):savename=self.savename
    if(data_form==None):data_form=self.data_form
    if(num_classes==None):num_classes=self.num_classes
    if(candidate==None):candidate=self.candidate
    if(channel==None):channel=self.channel
    if(extra==None):extra=self.extra
    if(balance==None):balance=self.balance
    if(len_train==None):len_train=self.len_train
    if(args==None):args=self.args
    print("##test$$")
    testX,testY,testE=self.getdata(data_form,num_classes,channel,extra,len_train,balance,args,energy)
    testp=self.model.predict(testX,verbose=0)
    fpr,tpr,thresholds=roc_curve(testY[:,0],testp[:,0])
    if(os.path.isfile("save/librarian.json"):
      book=json.load(open("save/librarian.json"))
      shutil.copyfile("save/librarian.json","save/backup/librarian{}.json".format(book['last']))
      index=book['last']
      for i in range(book['last']):
        if(book[str(i)]["savename"]==savename):
          index=i
          break
      if(index==book['last']):book['last']+=1
    else:
      book={'last':0}
      index=0
    book[str(index)]={"savename":savename,**vars(args),"AUC":roc_auc_score(testY[:,0],testp[:,0]),"start":str(self.start)}
    json.dump(book,open("save/librarian.json",'w'))

    tnr=1-fpr
    print("AUC:{}".format(round(roc_auc_score(testY[:,0],testp[:,0]),4)))
    np.savez("/pad/yulee/keras/drbox/{}out".format(args.save),y=testY,p=testp,testE=testE,trainy=self.trainY,trainp=self.trainp,trainE=self.trainE,classes=candidate)
    print("/pad/yulee/keras/drbox/{}out".format(args.save))
    del testX,testY
  def vis(self,savename=None,data_form=None,num_classes=None,candidate=None,balance=None,len_train=None,channel=None,args=None,num_conv=0):
    if(savename==None):savename=self.savename
    if(data_form==None):data_form=self.data_form
    if(num_classes==None):num_classes=self.num_classes
    if(candidate==None):candidate=self.candidate
    if(channel==None):channel=self.channel
    if(balance==None):balance=self.balance
    if(len_train==None):len_train=self.len_train
    if(args==None):args=self.args
    # summarize filter shapes
    count=0
    for layer in self.model.layers:
      if(count>num_conv):break
      # check for convolutional layer
      if 'conv' not in layer.name:
        continue
      # get filter weights
      filters, biases = layer.get_weights()
      print(layer.name, filters.shape)
      count+=1
    # normalize filter values to 0-1 so we can visualize them
    f_min, f_max = filters.min(), filters.max()
    filters = (filters - f_min) / (f_max - f_min)
    print(filters.shape)
    print("!@#!@#")
    # plot first few filters
    n_filters, ix = 2, 1
    fig=plt.figure()
    for i in range(n_filters):
      # get the filter
      f = filters[:, :, i]
      # plot each channel separately
      for j in range(3):
        # specify subplot and turn of axis
        ax = plt.subplot(n_filters, 3, ix)
        ax.set_xticks([])
        ax.set_yticks([])
        # plot filter channel in grayscale
        plt.imshow(f[:, :, j], cmap='gray')
        ix += 1
    # show the figure
    plt.show()
    return fig

if __name__ == '__main__':
  parser=argparse.ArgumentParser()
  parser.add_argument("--end",type=float,default=1.,help='end ratio')
  parser.add_argument("--save",type=str,default="test",help='save name model will be save in save/<argument>/')
  parser.add_argument("--cand",type=str,default="electron,pion",help='particle candidates split by ,')
  parser.add_argument("--network",type=str,default="0",help='network name on symbols/')
  parser.add_argument("--opt",type=str,default="sgd",help='optimizer sgd rms adam')
  parser.add_argument("--lr",type=float,default=0.000002,help='learning rate')
  parser.add_argument("--pt",type=str,default='',help='pt for dataset')
  parser.add_argument("--energy",type=str,default='',help='pt for dataset')
  parser.add_argument("--epochs",type=int,default=10,help='number of training epochs')
  parser.add_argument("--batch_size",type=int,default=128,help='batch_size')
  parser.add_argument("--loss",type=str,default="categorical_crossentropy",help='loss function')
  parser.add_argument("--gpu",type=int,default=0,help='0 < number set gpu number which uses, -1 == uses cpu')
  parser.add_argument("--pix",type=int,default=168,help='pixel size for image')
  parser.add_argument("--stride",type=str,default="label",help='end ratio')
  parser.add_argument("--target",type=int,default=0,help='target variable')
  parser.add_argument("--num_point",type=int,default=2048,help='number of points in point dataset')
  parser.add_argument("--channel",type=str,default="0,1,2,3",help='input channels')
  parser.add_argument("--extra",type=str,default="",help='input extra channels')
  parser.add_argument("--dform",type=str,default="point",help='point or image')
  parser.add_argument("--normalize",type=int,default=0,help='normalize at each index each channel')
  parser.add_argument("--mod",type=int,default=0,help='end ratio')
  parser.add_argument("--version",type=str,default="0.0.2.3",help='repo version')
  parser.add_argument("--pool",type=int,default=1,help='pixel pool size')
  parser.add_argument("--seed",type=int,default=123,help='random seed')
  parser.add_argument("--memo",type=str,default="",help='some memo')
  parser.add_argument("--job",type=str,default="train",help='some memo')
  args=parser.parse_args()
  savename='save/'+str(args.save)
  seed(args.seed)
  if(args.memo==""):args.memo=args.save
  if(not os.path.isdir(savename)):
    os.makedirs(savename)
  logging.basicConfig(filename=savename+'/log.log',level=logging.DEBUG)
  run1=drrun(args)
  run1.premodel()
  run1.train()
  run1.loadmodel()
  if(args.energy==''):energy=[-1,1000000]
  elif('-' in args.energy):energy=[int(i) for i in args.energy.split("-")]
  if(args.job=="train" or args.job=="test"):
    run1.test(energy=energy)
  if(args.job=="vis"):
    run1.vis()