#test_O, test_Rs, test_Rr, test_Ra, test_y = get_inputs(graphs[train_size:])
#if (verbose): print(" ... train_size={0}, test_size={1}".format(train_size, len(test_y)))

save_epochs = np.arange(0, n_epoch, save_every)
if (verbose):
    print(" ... saving the following epochs:", save_epochs)
    print(" ... saving the last {0} epochs".format(save_last))

test_losses, train_losses, batch_losses = [], [], []
for epoch in range(n_epoch):
    start_time = time.time()
    epoch_loss = 0.0
    batch_losses = []
    for b in range(n_batch):
        rand_idx = [random.randint(0, train_size) for _ in range(batch_size)]
        batch = [load_graph(graphs[i]) for i in rand_idx]
        O, Rs, Rr, Ra, y = get_inputs(batch)
        predicted = interaction_network(O, Rs, Rr, Ra)
        batch_loss = criterion(torch.cat(predicted, dim=0), torch.cat(y,
                                                                      dim=0))
        epoch_loss += batch_loss.item()
        batch_losses.append(batch_loss.item())
        optimizer.zero_grad()
        batch_loss.backward()
        optimizer.step()
        print(" ... batch {0}".format(b))

    end_time = time.time()
    train_losses.append(epoch_loss / n_batch)

    # compute test loss
                        for f in graph_files[train_size:]])
size = len(test_graphs)

# grab the IN with the best test performance
criterion = nn.BCELoss()
interaction_network = InteractionNetwork(3, 1, 1)
losses, epochs = [], []
for i, model in enumerate(models):
    epoch = int(model.split('epoch')[1].split('.')[0])
    #    if (epoch != 20): continue
    epochs.append(epoch)
    interaction_network.load_state_dict(torch.load(model))
    interaction_network.eval()
    model_loss_sum = 0
    for graph in test_graphs[:,1]:
        graph = load_graph(graph)
        O = [Variable(torch.FloatTensor(graph.X))]
        Rs = [Variable(torch.FloatTensor(graph.Ro))]
        Rr = [Variable(torch.FloatTensor(graph.Ri))]
        Ra = [Variable(torch.FloatTensor(graph.a)).unsqueeze(0)]
        y = [Variable(torch.FloatTensor(graph.y)).unsqueeze(0).t()]
        predicted = interaction_network(O, Rs, Rr, Ra)
        predicted = torch.cat(predicted, dim=0)
        target = torch.cat(y, dim=0)
        loss = (criterion(predicted, target))
        model_loss_sum += loss.item()
    
    model_loss = model_loss_sum/test_size
    losses.append(model_loss)
    print("Model {0} (Epoch #{1}):".format(model, i), model_loss)
tag = config['tag']
graph_indir = config['graph_indir']
train_size, test_size = 800, 200
plot_dir = "/tigress/jdezoort/IN_output/plots"

# job name, ex. "LP_0p5_1600_noPhi"
job_name = "{0}_{1}_{2}_{3}".format(prep, pt_cut, n_epoch, tag)
print("job_name={0}".format(job_name))

# load in test graph paths
graph_dir = "{0}/IN_{1}_{2}/".format(graph_indir, prep, pt_cut)
graph_files = os.listdir(graph_dir)

# (<evt_id>, <Graph>), e.g. evt_id=1582
test_graphs = np.array([(int(f.split('_')[0].split('t00000')[1]),
                         load_graph(graph_dir + f))
                        for f in graph_files[train_size:]])

# truth info, particle_id for each hit
particle_ids = [
    pickle.load(open("truth/" + str(test_graphs[i][0]) + "truth.pkl", 'rb'))
    for i in range(test_size)
]

# load in test_graphs
test_O = [
    Variable(torch.FloatTensor(test_graphs[i][1].X)) for i in range(test_size)
]
test_Rs = [
    Variable(torch.FloatTensor(test_graphs[i][1].Ro)) for i in range(test_size)
]
Esempio n. 4
0
def get_graphs(d):
    """ get all graphs from a directory
          return [("event<id1>", graph1, size), ...]
    """
    files = os.listdir(d)
    return [load_graph(d + f) for f in files]
              'n_segs'    : [], 'n_segs_er'    : [],
              'n_nodes'   : [], 'n_nodes_er'   : [],
              'size'      : [], 'size_er'      : []})

# analyze efficiencies per pt cut
for i in range(len(pt_cuts)):
    graph_files = os.listdir(data_dirs[i])
    seg_effs, truth_effs = [], []
    node_counts, edge_counts = [], []
    sizes = []

    # must loop over graphs due to memory limitations
    for graph_file in graph_files:
        graph_path = data_dirs[i] + graph_file
        graph_data = (graph_file.split('_')[0],
                      load_graph(graph_path),
                      os.path.getsize(graph_path)/10**6)
    
        sizes.append((sys.getsizeof(graph_data[1].X) + 
                     sys.getsizeof(graph_data[1].Ri) + 
                     sys.getsizeof(graph_data[1].Ro) + 
                     sys.getsizeof(graph_data[1].y) + 
                     sys.getsizeof(graph_data[1].a))/10.**6)

        n_nodes, n_edges = get_shape(graph_data)
        node_counts.append(n_nodes)
        edge_counts.append(n_edges)
        
        truth_eff = get_truth_efficiency(i, graph_data, truth_table)
        seg_eff = get_segment_efficiency(graph_data)
        truth_effs.append(truth_eff)