Merge branch 'master' of https://github.com/eric-moreno/IN

IN_dataGenerator.py

@@ -147,9 +147,7 @@ def main(args):
                    De=args.De,
                    Do=args.Do)
     else: 
-        gnn = GraphNetnoSV(N, n_targets, len(params), args.hidden, N_sv, len(params_sv),
-                       sv_branch=int(sv_branch),
-                       vv_branch=int(vv_branch),
+        gnn = GraphNetnoSV(N, n_targets, len(params), args.hidden,
                        De=args.De,
                        Do=args.Do)
 
@@ -178,6 +176,7 @@ def main(args):
 
     from sklearn.metrics import roc_curve, roc_auc_score, accuracy_score
     softmax = torch.nn.Softmax(dim=1)
+    import time
 
     for m in range(n_epochs):
         print("Epoch %s\n" % m)
@@ -187,7 +186,7 @@ def main(args):
         loss_val = []
         loss_training = []
         correct = []
-
+        tic = time.perf_counter()
         for sub_X,sub_Y,sub_Z in tqdm.tqdm(data_train.generate_data(),total=n_train/batch_size):
             training = sub_X[2]
             #training_neu = sub_X[1]
@@ -213,7 +212,9 @@ def main(args):
             optimizer.step()
             loss_string = "Loss: %s" % "{0:.5f}".format(l.item())
             del trainingv, trainingv_sv, targetv
-
+        toc = time.perf_counter()
+        print(f"Training done in {toc - tic:0.4f} seconds")
+        tic = time.perf_counter()
         for sub_X,sub_Y,sub_Z in tqdm.tqdm(data_val.generate_data(),total=n_val/batch_size):
             training = sub_X[2]
             #training_neu = sub_X[1]
@@ -239,7 +240,8 @@ def main(args):
 
             correct.append(target)
             del trainingv, trainingv_sv, targetv
-
+        toc = time.perf_counter()
+        print(f"Evaluation done in {toc - tic:0.4f} seconds")
         l_val = np.mean(np.array(loss_val))
 
         predicted = np.concatenate(lst) #(torch.FloatTensor(np.concatenate(lst))).to(device)

IN_onnx.py

@@ -3,14 +3,16 @@
 import setGPU
 import argparse
 import onnx
+import onnxruntime as ort
 import warnings
-#from onnx_tf.backend import prepare
 import os
-
+import sys
+import time
+
 N = 60 # number of charged particles
 N_sv = 5 # number of SVs 
 n_targets = 2 # number of classes
-save_path = '/bigdata/shared/BumbleB/convert_20181121_ak8_80x_deepDoubleB_db_pf_cpf_sv_dl4jets_test/'
+save_path = '/storage/group/gpu/bigdata/BumbleB/convert_20181121_ak8_80x_deepDoubleB_db_pf_cpf_sv_dl4jets_test/'
 
 params_2 = ['track_ptrel',     
           'track_erel',     
@@ -79,37 +81,56 @@ def main(args):
         gnn = GraphNet(N, n_targets, len(params), args.hidden, N_sv, len(params_sv),
                        vv_branch=int(args.vv_branch),
                        De=args.De,
-                       Do=args.Do)
+                       Do=args.Do,
+                       softmax=True)
 
     else:
-        gnn = GraphNetnoSV(N, n_targets, len(params), args.hidden, N_sv, len(params_sv),
-                           sv_branch=int(args.sv_branch),
-                           vv_branch=int(args.vv_branch),
+        gnn = GraphNetnoSV(N, n_targets, len(params), args.hidden, 
                            De=args.De,
-                           Do=args.Do)
+                           Do=args.Do,
+                           softmax=True)
 
-    gnn.load_state_dict(torch.load('%s/gnn_%s_best.pth'%(args.outdir,label)))
 
+    gnn.load_state_dict(torch.load('%s/gnn_%s_best.pth'%(args.outdir,label), map_location=torch.device('cpu')))
+
+    torch.save(gnn.state_dict(),'%s/gnn_%s_best_onnx.pth'%(args.outdir,label))
+
+    print(gnn)
     batch_size = 1
     dummy_input_1 = torch.from_numpy(test[0:batch_size]).cuda()
     dummy_input_2 = torch.from_numpy(test_sv[0:batch_size]).cuda()
     #dummy_input_1 = torch.randn(32, 30, 60, device='cuda')
     #dummy_input_2 = torch.randn(32, 14, 5, device='cuda')
 
     if args.sv_branch: 
+        tic = time.perf_counter()
         out_test = gnn(dummy_input_1, dummy_input_2)
+        toc = time.perf_counter()
         input_names = ['input_cpf', 'input_sv']
         output_names = ['output1']
         torch.onnx.export(gnn, (dummy_input_1, dummy_input_2), "%s/gnn.onnx"%args.outdir, verbose=True,
-                          input_names = input_names, output_names = output_names)
+                          input_names = input_names, output_names = output_names,
+                          export_params=True,        # store the trained parameter weights inside the model file
+                          opset_version=11,          # the ONNX version to export the model to
+                          dynamic_axes = {input_names[0]: {0: 'batch_size'}, 
+                                          input_names[1]: {0: 'batch_size'}, 
+                                          output_names[0]: {0: 'batch_size'}})
+
 
     else: 
+        tic = time.perf_counter()
         out_test = gnn(dummy_input_1)
+        toc = time.perf_counter()
         input_names = ['input_cpf']
         output_names = ['output1']
         torch.onnx.export(gnn, (dummy_input_1), "%s/gnn.onnx"%args.outdir, verbose=True,
-                          input_names = input_names, output_names = output_names)
+                          input_names = input_names, output_names = output_names,
+                          export_params=True,        # store the trained parameter weights inside the model file
+                          opset_version=11,          # the ONNX version to export the model to
+                          dynamic_axes = {input_names[0]: {0: 'batch_size'}, 
+                                          output_names[0]: {0, 'batch_size'}})
 
+    print(f"PyTorch Inference in {toc - tic:0.4f} seconds")
 
     # Load the ONNX model
     model = onnx.load("%s/gnn.onnx"%args.outdir)
@@ -120,6 +141,29 @@ def main(args):
     # Print a human readable representation of the graph
     print(onnx.helper.printable_graph(model.graph))
 
+    ort_session = ort.InferenceSession("%s/gnn.onnx"%args.outdir)
+
+    def to_numpy(tensor):
+        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
+
+    # compute ONNX Runtime output prediction
+    ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(dummy_input_1),
+                  ort_session.get_inputs()[1].name: to_numpy(dummy_input_2)}
+
+    tic = time.perf_counter()
+    ort_outs = ort_session.run(None, ort_inputs)
+    toc = time.perf_counter()
+    print(f"ONNXRuntime Inference in {toc - tic:0.4f} seconds")
+
+    print('PyTorch:', out_test)
+    print('ONNXRuntime:', ort_outs)
+    # compare ONNX Runtime and PyTorch results
+    np.testing.assert_allclose(to_numpy(out_test), ort_outs[0], rtol=1e-03, atol=1e-05)
+
+    print("Exported model has been tested with ONNXRuntime, and the result looks good!")
+
+
+    from onnx_tf.backend import prepare
     #warnings.filterwarnings('ignore') # Ignore all the warning messages in this tutorial
     tf_rep = prepare(model) # Import the ONNX model to Tensorflow
     print(tf_rep.inputs) # Input nodes to the model
@@ -133,7 +177,6 @@ def main(args):
     import tensorflow as tf
     tf.reset_default_graph()
     from tensorflow.python.platform import gfile
-    constDir = './'
     with tf.Session() as sess:
         with gfile.FastGFile(model_filename, 'rb') as f:
             graph_def = tf.GraphDef()
@@ -143,32 +186,15 @@ def main(args):
 
             x = sess.graph.get_tensor_by_name('import/input_cpf:0')
             y = sess.graph.get_tensor_by_name('import/input_sv:0')
-            out = sess.graph.get_tensor_by_name('import/add_33:0')
+            #out = sess.graph.get_tensor_by_name('import/add_33:0')
+            #out = sess.graph.get_tensor_by_name('import/Softmax_130:0')
+            out = sess.graph.get_tensor_by_name('import/output1:0')
             feed_dict = {x:test[0:batch_size], y:test_sv[0:batch_size]}
-
             classification = sess.run(out, feed_dict)
 
-            # make variables constant (seems to be unnecessary)
-
-            #from tensorflow.python.framework import graph_util
-            #from tensorflow.python.framework import graph_io
-
-            #pred_node_names = ['import/add_33']
-            #nonconstant_graph = tfsession.graph.as_graph_def()
-            #constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(), pred_node_names)
-
-            #f = 'constantgraph.pb.ascii'
-            #tf.train.write_graph(constant_graph, './', f, as_text=True)
-            #print('saved the graph definition in ascii format at: ', os.path.join('./', f))
-
-            #f = 'constantgraph.pb'
-            #tf.train.write_graph(constant_graph, './', f, as_text=False)
-            #print('saved the graph definition in pb format at: ', os.path.join('./', f))
-
-
-    print("pytorch\n",out_test)
-    #print("tensorflow\n",output)
-    print("tensorflow\n",classification)
+            sess.run(tf.global_variables_initializer())                                                                              
+    print("PyTorch:",out_test)
+    print("TensorFlow:",classification)
 
 
 if __name__ == "__main__":
@@ -181,9 +207,9 @@ def main(args):
     parser.add_argument("vv_branch", help="Required positional argument")
 
     # Optional arguments
-    parser.add_argument("--De", type=int, action='store', dest='De', default = 5, help="De")
-    parser.add_argument("--Do", type=int, action='store', dest='Do', default = 6, help="Do")
-    parser.add_argument("--hidden", type=int, action='store', dest='hidden', default = 15, help="hidden")
+    parser.add_argument("--De", type=int, action='store', dest='De', default = 20, help="De")
+    parser.add_argument("--Do", type=int, action='store', dest='Do', default = 24, help="Do")
+    parser.add_argument("--hidden", type=int, action='store', dest='hidden', default = 60, help="hidden")
 
     args = parser.parse_args()
     main(args)

data.py

@@ -118,6 +118,22 @@ def inf_generate_data_keras(self):
                     yield [B[2].swapaxes(1,2), B[3].swapaxes(1,2)], C
             except StopIteration as si:
                 print ("start over generator loop")
+
+    def inf_generate_data_keras_db(self):
+        while True:
+            try:
+                for B, C, D in self.generate_data():
+                    yield [B[0], B[2].swapaxes(1,2)[:,:,22:], B[3].swapaxes(1,2)[:,:,11:13]], C
+            except StopIteration as si:
+                print ("start over generator loop")
+
+    def inf_generate_data_keras_cpf(self):
+        while True:
+            try:
+                for B, C, D in self.generate_data():
+                    yield [B[2].swapaxes(1,2)], C
+            except StopIteration as si:
+                print ("start over generator loop")
 
     def generate_data(self):
        """Yields batches of training data until none are left."""

ddb.py

@@ -1,4 +1,4 @@
-from keras.layers import BatchNormalization, Conv1D, SpatialDropout1D, Dropout, Concatenate, Dense, Input, GRU
+from keras.layers import BatchNormalization, Conv1D, SpatialDropout1D, Dropout, Concatenate, Dense, Input, GRU, Flatten
 from keras.models import Model
 from keras.layers.advanced_activations import LeakyReLU
 
@@ -32,30 +32,33 @@ def model_DeepDoubleXReference(inputs, num_classes, datasets = ['cpf','sv'], hlf
 
     normalizedInputs = []
     for i in range(len(inputs)):
-        normedLayer = BatchNormalization(momentum=0.3,name = '%s_input_batchnorm'%datasets[i])(inputs[i])
+        normedLayer = BatchNormalization(momentum=0.3,axis=-1, name = '%s_input_batchnorm'%datasets[i])(inputs[i])
         normalizedInputs.append(normedLayer)
 
     flattenLayers = []
     if hlf_input is not None:
-        flattenLayers.append(normalizedInputs[hlf_input])
+        flattenLayers.append(Flatten()(normalizedInputs[hlf_input]))
 
     for i in range(len(inputs)):
         if i==hlf_input: continue
         ds = datasets[i]
         normalizedInput = normalizedInputs[i]
         x = Conv1D(filters=32*scale_hidden, kernel_size=(1,), strides=(1,), padding='same', 
                                 kernel_initializer=kernel_initializer, use_bias=False, name='%s_conv1'%ds, 
-                                activation = 'relu')(normalizedInput)
+                                activation = 'relu', data_format='channels_last')(normalizedInput)
         x = SpatialDropout1D(rate=0.1)(x)
         x = Conv1D(filters=32*scale_hidden, kernel_size=(1,), strides=(1,), padding='same', 
                              kernel_initializer=kernel_initializer, use_bias=False, name='%s_conv2'%ds, 
-                             activation = 'relu')(x)
+                             activation = 'relu', data_format='channels_last')(x)
         x = SpatialDropout1D(rate=0.1)(x)
         x = GRU(50*scale_hidden,go_backwards=True,implementation=2,name='%s_gru'%ds)(x)
         x = Dropout(rate=0.1)(x)
         flattenLayers.append(x)
 
-    concat = Concatenate()(flattenLayers)
+    if len(datasets)>1: 
+        concat = Concatenate()(flattenLayers)
+    else: 
+        concat = flattenLayers[0]
 
     fc = FC(concat, 100*scale_hidden, p=0.1, name='fc1')
     output = Dense(num_classes, activation='softmax', name='ID_pred', kernel_initializer=kernel_initializer_fc)(fc)

gnn.py

@@ -5,7 +5,7 @@
 import torch.optim as optim
 
 class GraphNet(nn.Module):
-    def __init__(self, n_constituents, n_targets, params, hidden, n_vertices, params_v, vv_branch=False, De=5, Do=6):
+    def __init__(self, n_constituents, n_targets, params, hidden, n_vertices, params_v, vv_branch=False, De=5, Do=6, softmax=False):
         super(GraphNet, self).__init__()
         self.hidden = int(hidden)
         self.P = params
@@ -23,6 +23,7 @@ def __init__(self, n_constituents, n_targets, params, hidden, n_vertices, params
         self.assign_matrices()
         self.assign_matrices_SV()
         self.vv_branch = vv_branch
+        self.softmax = softmax
         if self.vv_branch:
             self.assign_matrices_SVSV()
 
@@ -168,6 +169,9 @@ def forward(self, x, y):
         else:
             N = self.fc_fixed(N)
 
+        if self.softmax:
+            N = nn.Softmax(dim=-1)(N)
+
         return N 
 
     def tmul(self, x, y):  #Takes (I * J * K)(K * L) -> I * J * L 
@@ -262,13 +266,11 @@ def forward(self, x, y):
 
 # Architecture that excludes Secondary Vertices branch from Interaction network
 class GraphNetnoSV(nn.Module):
-    def __init__(self, n_constituents, n_targets, params, hidden, n_vertices, params_v, sv_branch=True, vv_branch=False, De=5, Do=6):
+    def __init__(self, n_constituents, n_targets, params, hidden, De=5, Do=6, softmax=False):
         super(GraphNetnoSV, self).__init__()
         self.hidden = int(hidden)
         self.P = params
         self.N = n_constituents
-        self.S = params_v
-        self.Nv = n_vertices
         self.Nr = self.N * (self.N - 1)
         self.Nt = self.N * self.Nv
         self.Ns = self.Nv * (self.Nv - 1)
@@ -278,8 +280,6 @@ def __init__(self, n_constituents, n_targets, params, hidden, n_vertices, params
         self.Do = Do
         self.n_targets = n_targets
         self.assign_matrices()
-        self.vv_branch = vv_branch
-        self.sv_branch = sv_branch
 
         self.Ra = torch.ones(self.Dr, self.Nr)
         self.fr1 = nn.Linear(2 * self.P + self.Dr, self.hidden).cuda()
@@ -294,7 +294,6 @@ def __init__(self, n_constituents, n_targets, params, hidden, n_vertices, params
         self.fo3 = nn.Linear(int(self.hidden/2), self.Do).cuda()
 
         self.fc_fixed = nn.Linear(self.Do, self.n_targets).cuda()
-        #self.gru = nn.GRU(input_size = self.Do, hidden_size = 20, bidirectional = False).cuda()
 
     def assign_matrices(self):
         self.Rr = torch.zeros(self.N, self.Nr)
@@ -343,6 +342,9 @@ def forward(self, x):
 
         N = self.fc_fixed(N)
 
+        if softmax:
+            N = nn.Softmax(dim=-1)(N)
+
         return N 
 
     def tmul(self, x, y):  #Takes (I * J * K)(K * L) -> I * J * L