diff --git a/IN_dataGenerator.py b/IN_dataGenerator.py index 1f3eb5d..890fbb8 100644 --- a/IN_dataGenerator.py +++ b/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) diff --git a/IN_onnx.py b/IN_onnx.py index 1c1532d..52f7c12 100644 --- a/IN_onnx.py +++ b/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,17 +81,21 @@ 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() @@ -97,19 +103,34 @@ def main(args): #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) diff --git a/data.py b/data.py index 6c28915..a88fc71 100644 --- a/data.py +++ b/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.""" diff --git a/ddb.py b/ddb.py index c836129..c3a90ae 100644 --- a/ddb.py +++ b/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,12 +32,12 @@ 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 @@ -45,17 +45,20 @@ def model_DeepDoubleXReference(inputs, num_classes, datasets = ['cpf','sv'], hlf 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) diff --git a/gnn.py b/gnn.py index e506bea..d5b02e8 100644 --- a/gnn.py +++ b/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 diff --git a/trained_models/IN.onnx b/trained_models/IN.onnx index 350e19f..7e77af0 100644 Binary files a/trained_models/IN.onnx and b/trained_models/IN.onnx differ diff --git a/trained_models/IN.pb b/trained_models/IN.pb index 652faf7..aea92c8 100644 Binary files a/trained_models/IN.pb and b/trained_models/IN.pb differ diff --git a/trained_models/IN.pth b/trained_models/IN.pth index 4da91de..e3f006d 100644 Binary files a/trained_models/IN.pth and b/trained_models/IN.pth differ