def train(model, num_epochs, dataset, train_loader, optimizer, criterion, device="cpu", k=5, all_loader=None): history = {'train': [], 'val': [], 'test': []} for epoch in range(num_epochs): model.train() for batch_size, n_id, adjs in train_loader: # `adjs` holds a list of `(edge_index, e_id, size)` tuples. adjs = [adj.to(device) for adj in adjs] num_users_per_res, users = dataset.get_visited_users( n_id[:batch_size], k) unique_users, inverse_idx = np.unique(users, return_inverse=True) unique_users = torch.LongTensor(unique_users) user_loader = NeighborSampler(dataset.user_pyg_graph.edge_index, node_idx=unique_users, sizes=[-1, -1], batch_size=unique_users.shape[0], shuffle=False) for num_users, u_id, u_adjs in user_loader: # this is actually just one loop: using for loop since # next(iter(user_loader)) appears to be buggy u_adjs = [adj.to(device) for adj in u_adjs] del user_loader optimizer.zero_grad() out = model(dataset.res_x[n_id], adjs, dataset.user_x[u_id], u_adjs, inverse_idx, num_users_per_res) # out = model(dataset.res_x[n_id], adjs) label = dataset.labels[n_id[:batch_size]] loss = criterion(out, label) loss.backward() optimizer.step() if all_loader is not None: train_acc, val_acc, test_acc = test(model, dataset, all_loader, device, k=k) print( f'epoch: {epoch + 1}, Train: {train_acc:.4f}, Val: {val_acc:.4f}, ' f'Test: {test_acc:.4f}') history['train'].append(train_acc) history['val'].append(val_acc) history['test'].append(test_acc) return history
def main(): dataset = 'ogb' device = 'cuda' if torch.cuda.is_available() else 'cpu' device = torch.device(device) num_layers = 3 sizes = [10 for _ in range(num_layers)] hidden_channels = 256 epochs = 200 lr = 0.01 log_steps = 1 batch_size = 1024 * 4 data = Dataset(root=Path('../../dataset'), name=dataset) train_loader = NeighborSampler(data.adj_t, sizes=sizes, node_idx=data.split_idx['train'], batch_size=batch_size, shuffle=True) subgraph_loader = NeighborSampler(data.adj_t, sizes=[-1, -1, -1], node_idx=None, batch_size=batch_size, shuffle=False) model = Graphsage(num_layers, data.x.shape[1], hidden_channels, data.num_classes).to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=lr) test_scores = [] for epoch in range(epochs + 1): loss = train(model, data, train_loader, optimizer, device) result = test(model, data, subgraph_loader, device) if epoch % log_steps == 0: train_res, valid_res, test_res = result test_scores.append(test_res) print( f'Run: {1:02d}, ' f'Epoch: {epoch:02d}, ' f'Loss: {loss:.4f}, ' f'Metric: {data.metric}', f'Train: {100 * train_res:.2f}%, ' f'Valid: {100 * valid_res:.2f}% ' f'Test: {100 * test_res:.2f}%') print(f"Best test accuracy: {max(test_scores) * 100:.2f}%")
def test_dataloader(self): return NeighborSampler(self.data.adj_t, node_idx=self.data.test_mask, sizes=[25, 10], return_e_id=False, transform=self.convert_batch, batch_size=2048, num_workers=3, persistent_workers=True)
def test_dataloader(self, batch_size=32, transforms=None): return NeighborSampler( self.data.adj_t[self.split_idx["test"]], node_idx=self.split_idx["test"], sizes=[-1, 10], batch_size=batch_size, shuffle=True, num_workers=self._num_workers, )
def train_dataloader(self): return NeighborSampler(self.adj_t, node_idx=self.train_idx, sizes=self.sizes, return_e_id=False, transform=self.convert_batch, batch_size=self.batch_size, shuffle=True, num_workers=4)
def train_model_and_save_embeddings(dataset, data, epochs, learning_rate, device): # Define Model encoder = EmbeddingEncoder(emb_dim=200, out_channels=64, n_nodes=dataset.num_nodes).to(device) decoder = CosineSimDecoder().to(device) model = VGAE(encoder=encoder, decoder=decoder).to(device) node_features, train_pos_edge_index = data.x.to( device), data.edge_index.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) # data.edge_index = data.edge_index.long() assert data.edge_index.max().item() < dataset.num_nodes data_loader = NeighborSampler(data, size=[25, 10], num_hops=2, batch_size=10000, shuffle=False, add_self_loops=False) model.train() for epoch in tqdm(range(epochs)): epoch_loss = 0.0 for data_flow in tqdm(data_loader()): optimizer.zero_grad() data_flow = data_flow.to(device) block = data_flow[0] embeddings = model.encode( node_features[block.n_id], block.edge_index ) # TODO Avoid computation of all node features! loss = model.recon_loss(embeddings, block.edge_index) loss = loss + (1 / len(block.n_id)) * model.kl_loss() epoch_loss += loss.item() # Compute gradients loss.backward() # Perform optimization step optimizer.step() z = model.encode(node_features, train_pos_edge_index) torch.save(z.cpu(), "large_emb.pt") print(f"Loss after epoch {epoch} / {epochs}: {epoch_loss}") return model
def get_neighbor_sampler(self, train_sizes, b1=50, b2=1000, nworkers=50): """Neighbor sampler for this data """ # neighbor sampler for training self.train_loader = NeighborSampler(self.edge_index, node_idx=None, sizes=train_sizes, batch_size=b1, num_workers=nworkers, shuffle=False) # neighbor sampler for testing (single-layer sampler) self.test_loader = NeighborSampler(self.edge_index, node_idx=None, sizes=[-1], batch_size=b2, num_workers=nworkers, shuffle=False)
def test_dataloader(self): # Test best validation model once again. return NeighborSampler( self.adj_t, node_idx=self.val_idx, sizes=self.sizes, return_e_id=False, transform=self.convert_batch, batch_size=self.batch_size, num_workers=2, )
def __init__(self, dataset, neighbor_sizes, node_types=None, metapaths=None, head_node_type=None, directed=True, resample_train=None, add_reverse_metapaths=True, inductive=False): self.neighbor_sizes = neighbor_sizes super(HeteroNeighborSampler, self).__init__(dataset, node_types, metapaths, head_node_type, directed, resample_train, add_reverse_metapaths, inductive) if self.use_reverse: self.add_reverse_edge_index(self.edge_index_dict) # Ensure head_node_type is first item in num_nodes_dict, since NeighborSampler.sample() function takes in index only the first num_nodes_dict = OrderedDict([(node_type, self.num_nodes_dict[node_type]) for node_type in self.node_types]) self.edge_index, self.edge_type, self.node_type, self.local_node_idx, self.local2global, self.key2int = \ group_hetero_graph(self.edge_index_dict, num_nodes_dict) self.int2node_type = { type_int: node_type for node_type, type_int in self.key2int.items() if node_type in self.node_types } self.int2edge_type = { type_int: edge_type for edge_type, type_int in self.key2int.items() if edge_type in self.edge_index_dict } self.neighbor_sampler = NeighborSampler(self.edge_index, node_idx=self.training_idx, sizes=self.neighbor_sizes, batch_size=128, shuffle=True)
def hidden_test_dataloader(self): return NeighborSampler( self.adj_t, node_idx=self.test_idx, sizes=self.sizes, return_e_id=False, transform=self.convert_batch, batch_size=self.batch_size, num_workers=3, )
def train_dataloader(self): return NeighborSampler(self.data.adj_t, node_idx=self.data.train_mask, sizes=[25, 10], return_e_id=False, transform=self.convert_batch, batch_size=1024, shuffle=True, num_workers=6, persistent_workers=True)
def create_neighbor_sampler(self, batch_size=2, stage=None): return NeighborSampler( self.data.edge_index, # the nodes that should be considered for sampling. node_idx=getattr(self.data, f"{stage}_mask"), # -1 indicates all neighbors will be selected sizes=[self._num_layers, -1], num_workers=self._num_workers, drop_last=self._drop_last, pin_memory=self._pin_memory, )
def forward(model, i, data): # 改成用 geometric的Data格式 items, targets, mask, batch, seq = data.x, data.y, data.sequence_mask, data.batch, data.sequence seq = seq.view(targets.shape[0], -1) mask = mask.view(targets.shape[0], -1) A = [] # datas = data.to_data_list() # graphs = [to_networkx(d) for d in datas] # A = [nx.convert_matrix.to_pandas_adjacency(g).values for g in graphs] # 無向圖adj = in + out # A_out = [g for g in graphs] # 有向圖的adj就是A_out # todo 解決cpu usage高的問題 # global graph gg = model.global_data gg_edge_index = gg.edge_index # 直接對 batch下所有node做NeighborSample batch_nodes = seq.flatten() # batch_nodes = torch.unique(batch_nodes) # 取unique node in batch sessions # batch_nodes = batch_nodes[batch_nodes!=0] # 移除padding node id # sample as whole batch, 從大graph中找session graph內的node id的鄰居 # subgraph_loaders = NeighborSampler(gg_edge_index, node_idx=batch_nodes, sizes=[-1], shuffle=False, num_workers=0, batch_size=batch_nodes.shape[0]) # all neighbors # fixme 放全部node subgraph_loaders = NeighborSampler( gg_edge_index, node_idx=batch_nodes, sizes=[10, 5], shuffle=False, num_workers=0, batch_size=batch_nodes.shape[0]) # 2 hop hidden, pad, g_h = model(items, A, data.edge_index, subgraph_loaders) # session graph node embeddings # 推回原始序列 sections = torch.bincount(batch).cpu().numpy() # split whole x back into graphs G_i hidden = torch.split(hidden, tuple(sections)) # todo 增加不考慮padding的選項 mask_true = True if mask_true: leng = mask.shape[1] # padding完的session長度 alias_inputs = data.alias_inputs s_len = data.sequence_len.cpu().numpy().tolist() alias_inputs = torch.split(alias_inputs, s_len) seq_hidden = torch.stack([ get(pad, i, hidden, alias_inputs, leng) for i in torch.arange(len(alias_inputs)).long() ]) g_h = g_h.view([len(hidden), leng, -1]) else: seq_hidden = hidden seq_hidden += g_h return targets, model.compute_scores(seq_hidden, mask, mask_true)
def make_graph_sampler(self): graph = Data( edge_index=self.edge_index, edge_attr=self.edge_weight, num_nodes=self.num_nodes ) .to('cpu') graph_sampler = NeighborSampler( # graph, size=[5, 5], num_hops=2, batch_size=100, shuffle=self.shuffle, add_self_loops=True graph, size=[10, 15], num_hops=2, batch_size=250, shuffle=self.shuffle, add_self_loops=True ) return graph_sampler
def fit(self, model, data): data = data[0] self.model = model.to(self.device) self.data = data self.test_gpu_volume() self.subgraph_loader = NeighborSampler( data.edge_index, sizes=[ -1, ], batch_size=self.batch_size, shuffle=False, num_workers=10, ) self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=self.weight_decay) self.edge_index, _ = add_remaining_self_loops( data.edge_index, torch.ones(data.edge_index.shape[1]).to(data.x.device), 1, data.x.shape[0]) self.train_index = torch.where(data.train_mask)[0].tolist() epoch_iter = tqdm(range(self.max_epoch)) patience = 0 best_score = 0 best_loss = np.inf max_score = 0 min_loss = np.inf for epoch in epoch_iter: self._train_step() if epoch % 5 == 0: val_acc, val_loss = self._test_step(split="val") epoch_iter.set_description( f"Epoch: {epoch:03d}, Val: {val_acc:.4f}") if val_loss <= min_loss or val_acc >= max_score: if val_acc >= best_score: # SAINT loss is not accurate best_loss = val_loss best_score = val_acc best_model = copy.deepcopy(self.model) min_loss = np.min((min_loss, val_loss)) max_score = np.max((max_score, val_acc)) patience = 0 else: patience += 1 if patience == self.patience: self.model = best_model epoch_iter.close() break test_acc, _ = self._test_step(split="test") val_acc, _ = self._test_step(split="val") print(f"Test accuracy = {test_acc}") return dict(Acc=test_acc, ValAcc=val_acc)
def create_neighbor_sampler(self, batch_size=2, transforms=None, stage=None): return NeighborSampler( self.data.edge_index, node_idx=getattr(self.data, f"{stage}_mask"), sizes=self._sizes, num_workers=self._num_workers, drop_last=self._drop_last, pin_memory=self._pin_memory, )
def create_loader_from_cls(self, loader_cls=None, params=None, stage=None): if not hasattr(self, "_loaded_dataset"): self._loaded_dataset = loader_cls(self.data, **params) dataset = Batch.from_data_list([d for d in self._loaded_dataset]) return NeighborSampler( dataset.edge_index, node_idx=getattr(dataset, f"{stage}_mask"), sizes=self._sizes, num_workers=self._num_workers, drop_last=self._drop_last, pin_memory=self._pin_memory, )
def create_neighbor_sampler(self, batch_size=2, stage=None): # https://github.com/rusty1s/pytorch_geometric/tree/master/torch_geometric/data/sampler.py#L18 return NeighborSampler( self.data.edge_index, # the nodes that should be considered for sampling. node_idx=getattr(self.data, f"{stage}_mask"), # -1 indicates all neighbors will be selected sizes=[self._num_layers, -1], num_workers=self._num_workers, drop_last=self._drop_last, pin_memory=self._pin_memory, )
def test_sampler(): num_nodes = 10 data = Data(edge_index=erdos_renyi_graph(num_nodes, 0.1)) data.num_nodes = num_nodes loader = NeighborSampler(data, size=[4, 0.5], num_hops=2, batch_size=2, shuffle=True) for data_flow in loader(): assert data_flow.__repr__()[:8] == 'DataFlow' assert data_flow.n_id.size() == (2, ) assert data_flow.batch_size == 2 assert len(data_flow) == 2 block = data_flow[0] assert block.__repr__()[:5] == 'Block' for block in data_flow: pass data_flow = data_flow.to(torch.long) break for data_flow in loader(torch.tensor([0, 1, 2, 3, 4])): pass loader = NeighborSampler(data, size=[4, 0.5], num_hops=2, batch_size=3, drop_last=True, shuffle=False, add_self_loops=True) for data_flow in loader(): pass for data_flow in loader(torch.tensor([0, 1, 2, 3, 4])): pass mask = torch.tensor([0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=torch.uint8) for data_flow in loader(mask): pass
def test_sampler(): torch.manual_seed(12345) edge_index = erdos_renyi_graph(num_nodes=10, edge_prob=0.5) E = edge_index.size(1) loader = NeighborSampler(edge_index, sizes=[2, 4], batch_size=2) assert loader.__repr__() == 'NeighborSampler(sizes=[2, 4])' assert len(loader) == 5 for batch_size, n_id, adjs in loader: assert batch_size == 2 assert all(np.isin(n_id, torch.arange(10)).tolist()) assert n_id.unique().size(0) == n_id.size(0) for (edge_index, e_id, size) in adjs: assert int(edge_index[0].max() + 1) <= size[0] assert int(edge_index[1].max() + 1) <= size[1] assert all(np.isin(e_id, torch.arange(E)).tolist()) assert e_id.unique().size(0) == e_id.size(0) assert size[0] >= size[1] out = loader.sample([1, 2]) assert len(out) == 3
def train(model, num_epochs, dataset, optimizer, criterion, device="cpu", k=5): for epoch in range(num_epochs): model.train() for target_bus, business_ids, distances, business_lens, labels in dataset.get_batch( 2048, "train"): num_users_per_res, users = dataset.get_visited_users(target_bus, k) unique_users, inverse_idx = np.unique(users, return_inverse=True) unique_users = torch.LongTensor(unique_users) user_loader = NeighborSampler(dataset.user_pyg_graph.edge_index, node_idx=unique_users, sizes=[-1, -1], batch_size=unique_users.shape[0], shuffle=False) for _, u_id, u_adjs in user_loader: # this is actually just one loop: using for loop since # next(iter(user_loader)) appears to be buggy u_adjs = [adj.to(device) for adj in u_adjs] del user_loader business_ids = business_ids.to(device) business_lens = business_lens.to(device) labels = labels.to(device) optimizer.zero_grad() out = model(dataset.res_x, target_bus, business_ids, distances, business_lens, dataset, device, user_x=dataset.user_x[u_id], u_adjs=u_adjs, inverse_idx=inverse_idx, num_users_per_res=num_users_per_res) loss = criterion(out, labels) loss.backward() optimizer.step() if epoch % 5 == 0: train_acc = test(model, dataset, device, "train", k=k) val_acc = test(model, dataset, device, "val", k=k) test_acc = test(model, dataset, device, "test", k=k) print( f'epoch: {epoch}, Train: {train_acc:.4f}, Val: {val_acc:.4f}, ' f'Test: {test_acc:.4f}')
def _run_trainer(): path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'Reddit') print("Load Dataset") dataset = Reddit(path) data = dataset[0] print("Load Train Sampler") train_loader = NeighborSampler(data.edge_index, node_idx=data.train_mask, sizes=[25, 10], batch_size=1024, shuffle=True, num_workers=12) print("Load subgraph sampler") subgraph_loader = NeighborSampler(data.edge_index, node_idx=None, sizes=[-1], batch_size=1024, shuffle=False, num_workers=12) print("Creating SAGE model") model = SAGE('ps', dataset.num_features, 256, dataset.num_classes, subgraph_loader) optimizer = DistributedOptimizer(torch.optim.Adam, model.parameter_rrefs(), lr=0.01) print("Start training") for epoch in range(1, 11): loss, acc = train(model, optimizer, epoch, data, train_loader) print(f'Epoch {epoch:02d}, Loss: {loss:.4f}, Approx. Train: {acc:.4f}') train_acc, val_acc, test_acc = test(model, data) print(f'Train: {train_acc:.4f}, Val: {val_acc:.4f}, ' f'Test: {test_acc:.4f}')
def test(model, dataset, device, mode, k=5): outs = [] model.eval() y_true = [] for target_bus, business_ids, distances, business_lens, labels in dataset.get_batch( 2048, mode): num_users_per_res, users = dataset.get_visited_users(target_bus, k) unique_users, inverse_idx = np.unique(users, return_inverse=True) unique_users = torch.LongTensor(unique_users) user_loader = NeighborSampler(dataset.user_pyg_graph.edge_index, node_idx=unique_users, sizes=[-1, -1], batch_size=unique_users.shape[0], shuffle=False) for _, u_id, u_adjs in user_loader: # this is actually just one loop: using for loop since # next(iter(user_loader)) appears to be buggy u_adjs = [adj.to(device) for adj in u_adjs] del user_loader business_ids = business_ids.to(device) business_lens = business_lens.to(device) out = model(dataset.res_x, target_bus, business_ids, distances, business_lens, dataset, device, user_x=dataset.user_x[u_id], u_adjs=u_adjs, inverse_idx=inverse_idx, num_users_per_res=num_users_per_res) outs.append(out.cpu()) y_true.append(labels.cpu()) outs = torch.cat(outs, dim=0) y_true = torch.cat(y_true, dim=0) y_pred = outs.argmax(dim=-1, keepdim=False) results = int(y_pred.eq(y_true).sum()) / int(y_true.shape[0]) return results
def train_net(): club = KarateClub() data = club.data data.num_nodes = data.num_nodes[0] print(data) data_loader = NeighborSampler(data, size=[20, 10], num_hops=2, batch_size=8, shuffle=True, add_self_loops=True) net = SAGENet(34, 2) criterion = nn.NLLLoss() optimizer = torch.optim.Adam(net.parameters(), lr=lr) losses = AverageMeter() for i in range(epoch): result = [] for data_flow in data_loader(): label = data.y[data_flow.n_id] # print(label) optimizer.zero_grad() out = net(data.x, data_flow.to(device)) # print(out) # print(label) result.append((label, out)) # Calculate loss loss = criterion(out, label) # Back prop. optimizer.zero_grad() loss.backward() # Clip gradients clip_gradient(optimizer, grad_clip) # Update weights optimizer.step() losses.update(loss.item()) # Print status print(validation(result)) if i % print_freq == 0: print('Epoch: {0} Loss {loss.val:.5f} ({loss.avg:.5f})\t'.format( i, loss=losses)) return losses.avg
def make_sampler(self, node_idx, force_default=False): if node_idx != None: node_idx = torch.tensor(node_idx, dtype=torch.long) sizes = config["sample_sizes"] else: sizes = [-1] if config["unsupervised_loss"] and not force_default: if config["unsup_sampling_type"] == 'neighbor': adj_t = SparseTensor(row=self.graph.edge_index[0], col=self.graph.edge_index[1], value=self.graph.edge_attr, sparse_sizes=(self.graph.num_nodes, self.graph.num_nodes)).t() return PosNegNeighborSampler(edge_index=adj_t, node_idx=node_idx, sizes=sizes, batch_size=config["sample_batch_size"], shuffle=True, num_doc_nodes=len(self.docs), edge_weights=self.graph.edge_attr) else: raise NotImplementedError("[dataset] The unsupervised sampling type %s has not been implemented" % config["unsup_sampling_type"]) return NeighborSampler(self.graph.edge_index, node_idx=node_idx, sizes=sizes, batch_size=config["sample_batch_size"], shuffle=True, num_workers=config["sampling_num_workers"])
def _run_trainer(): path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'Reddit') print("Load Dataset") dataset = Reddit(path) data = dataset[0] print("Load Train Sampler") train_loader = NeighborSampler(data.edge_index, node_idx=data.train_mask, sizes=[25, 10], batch_size=1024, shuffle=True, num_workers=0) print("Creating SAGE model") model = SAGE(dataset.num_features, 128, dataset.num_classes) print("Start training") for epoch in range(1, 11): loss, acc = train(model, epoch, data, train_loader) print(f'Epoch {epoch:02d}, Loss: {loss:.4f}, Approx. Train: {acc:.4f}')
def test(model, dataset, all_loader, device, k=5): outs = [] model.eval() for batch_size, n_id, adjs in all_loader: adjs = [adj.to(device) for adj in adjs] num_users_per_res, users = dataset.get_visited_users( n_id[:batch_size], k) unique_users, inverse_idx = np.unique(users, return_inverse=True) unique_users = torch.LongTensor(unique_users) user_loader = NeighborSampler(dataset.user_pyg_graph.edge_index, node_idx=unique_users, sizes=[-1, -1], batch_size=unique_users.shape[0], shuffle=False) for num_users, u_id, u_adjs in user_loader: # this is actually just one loop: using for loop since # next(iter(user_loader)) appears to be buggy u_adjs = [adj.to(device) for adj in u_adjs] del user_loader out = model(dataset.res_x[n_id], adjs, dataset.user_x[u_id], u_adjs, inverse_idx, num_users_per_res) # out = model(dataset.res_x[n_id], adjs) outs.append(out.cpu()) outs = torch.cat(outs, dim=0) y_true = dataset.labels.cpu().unsqueeze(-1) y_pred = outs.argmax(dim=-1, keepdim=True) results = [] for idx in [dataset.train_index, dataset.val_index, dataset.test_index]: results.append( int(y_pred[idx].eq(y_true[idx]).sum()) / int(idx.shape[0])) return results
def detect(): global model global loader global data global device global fp global tn global id_map_t loader = NeighborSampler(data, size=[1.0, 1.0], num_hops=2, batch_size=batch_size, shuffle=True, add_self_loops=True) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') Net = SAGENet test_acc = 0 model = Net(feature_num, label_num).to(device) for j in model_list: loop_num = 0 base_model = str(j) while (1): model_path = '../models/' + base_model + '_' + str(loop_num) if not osp.exists(model_path): break model.load_state_dict(torch.load(model_path)) fp = [] tn = [] loss, test_acc = test(data.test_mask) for i in tn: data.test_mask[i] = False update_benign(i) if test_acc == 1: break loop_num += 1 if test_acc == 1: break for i in fp: raise_alert(i)
def test_cora(): class Net(torch.nn.Module): def __init__(self, in_channels, out_channels): super(Net, self).__init__() self.conv1 = SAGEConv(in_channels, 16) self.conv2 = SAGEConv(16, 16) self.conv3 = SAGEConv(16, out_channels) def forward_data_flow(self, x, edge_weight, data_flow): block = data_flow[0] weight = edge_weight[block.e_id] weight[block.e_id == -1] = 1 x = relu(self.conv1(x, block.edge_index, weight, block.size)) block = data_flow[1] weight = edge_weight[block.e_id] weight[block.e_id == -1] = 1 x = relu(self.conv2(x, block.edge_index, weight, block.size)) block = data_flow[2] weight = edge_weight[block.e_id] weight[block.e_id == -1] = 1 x = self.conv3(x, block.edge_index, weight, block.size) return x def forward(self, x, edge_index, edge_weight): x = relu(self.conv1(x, edge_index, edge_weight)) x = relu(self.conv2(x, edge_index, edge_weight)) return self.conv3(x, edge_index, edge_weight) root = osp.join('/', 'tmp', str(random.randrange(sys.maxsize))) dataset = Planetoid(root, 'Cora') model = Net(dataset.num_features, dataset.num_classes) data1 = dataset[0] data1.edge_weight = torch.rand(data1.num_edges) data2 = T.AddSelfLoops()(dataset[0]) data2.edge_weight = torch.rand(data2.num_edges) data3 = dataset[0] loop = torch.stack([torch.arange(100, 200), torch.arange(100, 200)], dim=0) data3.edge_index = torch.cat([data3.edge_index, loop], dim=1) data3.edge_weight = torch.rand(data3.num_edges) for data in [data1, data2, data3]: out_all = model(data.x, data.edge_index, data.edge_weight) loader = NeighborSampler(data, size=1.0, num_hops=3, batch_size=64, shuffle=False, drop_last=False, bipartite=True, add_self_loops=True) for data_flow in loader(data.train_mask): out = model.forward_data_flow(data.x[data_flow[0].n_id], data.edge_weight, data_flow) assert torch.allclose(out_all[data_flow.n_id], out) loader = NeighborSampler(data, size=1.0, num_hops=3, batch_size=64, shuffle=False, drop_last=False, bipartite=False) for subdata in loader(data.train_mask): out = model(data.x[subdata.n_id], subdata.edge_index, data.edge_weight[subdata.e_id]) out = out[subdata.sub_b_id] assert torch.allclose(out_all[subdata.b_id], out) shutil.rmtree(root)
from torch.nn import Linear as Lin from tqdm import tqdm from ogb.nodeproppred import PygNodePropPredDataset, Evaluator from torch_geometric.data import NeighborSampler from torch_geometric.nn import GATConv root = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'products') dataset = PygNodePropPredDataset('ogbn-products', root) split_idx = dataset.get_idx_split() evaluator = Evaluator(name='ogbn-products') data = dataset[0] train_idx = split_idx['train'] train_loader = NeighborSampler(data.edge_index, node_idx=train_idx, sizes=[10, 10, 10], batch_size=512, shuffle=True, num_workers=12) subgraph_loader = NeighborSampler(data.edge_index, node_idx=None, sizes=[-1], batch_size=1024, shuffle=False, num_workers=12) class GAT(torch.nn.Module): def __init__(self, in_channels, hidden_channels, out_channels, num_layers, heads): super(GAT, self).__init__()