Problem with batchsize on edge_index

[Volodimirich]

Hi, i am working with GNN on images and have found a strange bug with edge_index. I can't store all dataset in memory, and because of it, I am working with Dataset class:

SIZE = 256
files = glob.glob('/cats_dogs_dataset/preprocessed/*/*.npy', recursive=True)
from torch_geometric.data import Dataset, download_url
from torch_geometric.data import Data
import os.path as osp

class MyOwnDataset(Dataset):
    def __init__(self, root, files_list, is_train, size=256, allow_loops = True, transform=None, pre_transform=None, pre_filter=None):
        self.data = files_list
        self.allow_loops = allow_loops
        self.is_train = is_train
        #TODO fix this 
        self.gl_count = 1750 if self.is_train else 750
        super().__init__(root, transform, pre_transform, pre_filter)

    @property
    def raw_file_names(self):
        return []
    
    @property
    def processed_file_names(self):
        return [f'data_{idx}_is_train_{self.is_train}.pt' for idx in range(self.gl_count)]

    def _create_cco_matrix(self, allow_loops=False):
        result = split_tensor(SIZE)
        vert_amount = int(result[2][-1][-1]+1)
        adj_matrix = np.zeros((vert_amount,vert_amount)).astype(int)
        
        gl_map, gen_map, loc_map = result[0].astype(int), result[1].astype(int), result[2].astype(int)
        for gl, gen, loc in np.nditer([gl_map, gen_map, loc_map]):
            adj_matrix[gl, gen] = 1
            adj_matrix[gl, loc] = 1
            adj_matrix[gen, loc] = 1
        

        source_nodes = []
        target_nodes = []
        edge_list = []
        for iy, ix in np.ndindex(adj_matrix.shape):
            if adj_matrix[iy, ix] == 1:
                if iy == ix and not allow_loops:
                    continue
                source_nodes.append(ix)
                target_nodes.append(iy)
                
                #unweighted solution
                edge_list.append(1)
                
        return source_nodes, target_nodes, edge_list

        

    def process(self):
        idx = 0
        source_vert, target_vert, edge_list = self._create_cco_matrix()
        edge_idx = torch.tensor([source_vert, target_vert])
        
        #DEBUG ROW
        _, small_data = train_test_split(self.data, test_size = 0.1, random_state=42)
        train_data, test_data = train_test_split(small_data, test_size = 0.3, random_state=42)
        data = train_data if self.is_train else test_data
        
        for file in data:
            amp, phase, target = np.load(file, allow_pickle = True)   

            #temporary only amp|
            data = Data(x = amp, 
                    edge_index = torch.tensor(edge_idx),
                    edge_attrs = edge_list, 
                    y = torch.tensor([target]))

            torch.save(data, osp.join(self.processed_dir, f'data_{idx}_is_train_{self.is_train}.pt'))
            idx += 1
        self.gl_count = idx

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        data = torch.load(osp.join(self.processed_dir, f'data_{idx}_is_train_{self.is_train}.pt'))
        print(data.edge_index)
        return data

And i have next model initilization:

save_root = '/cats_dogs_dataset/'
files = glob.glob('/cats_dogs_dataset/preprocessed/*/*.npy', recursive=True)
train_dataset = MyOwnDataset(save_root, files, is_train = True)
test_dataset = MyOwnDataset(save_root, files, is_train = False)

from torch_geometric.loader import DataLoader


model = GCN(num_classes=2, hidden_channels = 10).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr = 0.001, weight_decay = 0.001)
criterion = torch.nn.CrossEntropyLoss()

batch_size = 2
train_loader = DataLoader(train_dataset, batch_size = batch_size, shuffle = True)
test_loader = DataLoader(test_dataset, batch_size = batch_size, shuffle = False)

My train loop:

  def train():
      model.train()
      itt = 0
      for itt, data in enumerate(train_loader):  # Iterate in batches over the training dataset.
          print(data.edge_index)
          optimizer.zero_grad()  # Clear gradients.
          data = data.to(device)
          out = model(data.x, data.edge_index, data.batch)  # Perform a single forward pass.
          loss = criterion(out.squeeze(0), data.y)  # Compute the loss.
          print(loss)
          if itt % 100 == 0:
              print(loss)
          loss.backward()  # Derive gradients.
          optimizer.step()  # Update parameters based on gradients.

My GNN graph is static and i always have 896 nodes. And with batch size equal to 1 I have correct edge-index:

tensor([[191, 255, 319,  ..., 767, 863, 895],
        [  0,   0,   0,  ..., 383, 383, 383]])

Now I have next problem, as soon as I set batch_size != 1 my code is crushed due to edge_index:

RuntimeError: INDICES element is out of DATA bounds, id=1087 axis_dim=896

I found the problem after some debugging. For some reason, batch size increases edge_index values. For example, with batch_size equal to 2:

#Debug row from get method from my OwnDataset. Everything is correct
tensor([[191, 255, 319,  ..., 767, 863, 895],
        [  0,   0,   0,  ..., 383, 383, 383]])

#Debug row from get method from my OwnDataset. Everything is correct
tensor([[191, 255, 319,  ..., 767, 863, 895],
        [  0,   0,   0,  ..., 383, 383, 383]])

#Debug from train loop. As I assume, train_loader concatenated the edge_index and added 896 (the number of vertices) to all elements of the second copy
tensor([[ 191,  255,  319,  ..., 1663, 1759, 1791],
        [   0,    0,    0,  ..., 1279, 1279, 1279]])

Because of it i can't work with batch size > 1. Can you please tell me where I went wrong?

[rusty1s]

This is intended. You should reshape your features to x.view(batch_size * num_nodes, -1).

[Volodimirich]

Thank you, it worked!)