edge_index when batching

[sachinmloecher]

Hi everyone, I am struggling with getting the edge_index attribute right when microbatching.

I have a Torch Dataset with the following methods:

def graph_response(self, response: t.Union[np.ndarray, torch.Tensor]):
        node_features = self.transform_response(response)
        data = Data(x=node_features, edge_index=self.edge_index.t())
        return data

def __getitem__(self, idx: t.Union[int, torch.Tensor], include_behavior=True, include_pupil=True):
        """Return data and metadata

        Returns
            - If args.as_graph is False:
                - data, t.Dict[str, torch.Tensor]
                    - image: the natural image in (C, H, W) format
                    - response: the corresponding response
                    - behavior: pupil size, the derivative of pupil size, and speed
                    - pupil_center: the (x, y) coordinate of the center of the pupil
                    - image_id: the frame image ID
                    - trial_id: the trial ID
                    - mouse_id: the mouse ID
            - If args.as_graph is True:
                - data, torch_geometric.data.Data
                    - image: the natural image in (C, H, W) format
                    - response: the corresponding response as a graph (x, edge_index)
                    - behavior: pupil size, the derivative of pupil size, and speed
                    - pupil_center: the (x, y) coordinate of the center of the pupil
                    - image_id: the frame image ID
                    - trial_id: the trial ID
                    - mouse_id: the mouse ID
        """
        trial = self.indexes[idx]
        data = load_trial_data(mouse_dir=self.mouse_dir, trial=trial)
        data["image"] = self.transform_image(data["image"])
        
        if self.as_graph:
            response = self.graph_response(data["response"])
            response.image = data["image"]
            response.behavior = self.transform_behavior(data["behavior"])
            response.pupil_center = self.transform_pupil_center(data["pupil_center"])
            response.image_id = self.image_ids[idx]
            response.trial_id = self.trial_ids[idx]
            response.mouse_id = self.mouse_id
            return response
        else:
            data["response"] = self.transform_response(data["response"])
            data["behavior"] = self.transform_behavior(data["behavior"])
            data["pupil_center"] = self.transform_pupil_center(data["pupil_center"])
            data["image_id"] = self.image_ids[idx]
            data["trial_id"] = self.trial_ids[idx]
            data["mouse_id"] = self.mouse_id
        
            return data

I am using the torch_geometric.loader.DataLoader and am a bit confused as to how batching works. I understand that each batch of Data objects is merged into one large Data object, but are the batch edge_indexes also supposed to be merged into one large edge_index? Here is my microbatching function:

def micro_batching(batch: Data, micro_batch_size: int):
    num_samples = len(batch.image)
    num_edges = batch.edge_index.size(1)
    edges_per_sample = num_edges // num_samples
    for i in range(0, num_samples, micro_batch_size):
        start_edge = i * edges_per_sample
        end_edge = (i + micro_batch_size) * edges_per_sample
        micro_batch = Data(
            x=batch.x[i:i+micro_batch_size],
            edge_index=batch.edge_index[:, 0:edges_per_sample],
            # edge_index=batch.edge_index[:, start_edge:end_edge],
            image=batch.image[i:i+micro_batch_size],
            behavior=batch.behavior[i:i+micro_batch_size],
            pupil_center=batch.pupil_center[i:i+micro_batch_size],
            image_id=batch.image_id[i:i+micro_batch_size],
            trial_id=batch.trial_id[i:i+micro_batch_size],
            mouse_id=batch.mouse_id[i:i+micro_batch_size]
        )
        yield micro_batch

I am currently just taking the edge_index from 1 sample and using that for the whole batch (my edges are static across samples). In my train step function I then do:

def train_step(
    mouse_id: str,
    batch: Data,
    model: FullModel,
    optimizer: torch.optim,
    criterion: torch.nn.Module,
    update: bool,
    micro_batch_size: int,
    device: torch.device = "cpu"
):
    model.to(device)
    batch_size = len(batch.image)
    result = {"loss/loss": []}
    for micro_batch in data.micro_batching(batch, micro_batch_size):
        y_true = torch.stack([torch.from_numpy(img) for img in micro_batch.image]).to(device)
        print(micro_batch.edge_index)
        print(np.array(micro_batch.x).shape)
        y_pred = model(
            x=torch.stack([torch.from_numpy(x) for x in micro_batch.x]).to(device),
            edge_index=micro_batch.edge_index.to(device),
            # edge_index=torch.stack([edge_index for edge_index in micro_batch.edge_index]).to(device),
            mouse_id=mouse_id,
            behaviours=torch.stack([torch.from_numpy(behavior) for behavior in micro_batch.behavior]).to(device),
            pupil_centers=torch.stack([torch.from_numpy(pupil_center) for pupil_center in micro_batch.pupil_center]).to(device)
        )
...

And eventually get this error when using a GCNConv layer:

C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\cuda\ScatterGatherKernel.cu:145: block: [33,0,0], thread: [63,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && "index out of bounds"` failed.
Traceback (most recent call last):
  File "C:\Users\sachi\SensoriumDecoding\src\models\trainGCN.py", line 316, in <module>
    model = main(args)
            ^^^^^^^^^^
  File "C:\Users\sachi\SensoriumDecoding\src\models\trainGCN.py", line 280, in main
    train_result = train(
                   ^^^^^^
  File "C:\Users\sachi\SensoriumDecoding\src\models\trainGCN.py", line 169, in train
    result = train_step(
             ^^^^^^^^^^^
  File "C:\Users\sachi\SensoriumDecoding\src\models\trainGCN.py", line 129, in train_step
    y_pred = model(
             ^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\SensoriumDecoding\src\models\model.py", line 27, in forward
    mouse_output = self.mouse_mlp_dict[mouse_id](x, edge_index)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\SensoriumDecoding\src\models\mouseGCN.py", line 36, in forward
    x = gcn_layer(x, edge_index)
        ^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch_geometric\nn\conv\gcn_conv.py", line 241, in forward
    edge_index, edge_weight = gcn_norm(  # yapf: disable
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch_geometric\nn\conv\gcn_conv.py", line 109, in gcn_norm
    deg_inv_sqrt = deg.pow_(-0.5)
                   ^^^^^^^^^^^^^^
RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

I have made sure that my edge_index does not contain any self loops, and all edge indices are within bounds. The edge_index being passed to the model in the microbatch is of size (2, 25116) and microbatch.x is of size (16, 8372) where the microbatch size is 16.
I am clearly getting something wrong with the edge_index during microbatching but am not sure how it is usually done in PyG. Thanks for any help!

[rusty1s]

Yes, DataLoader takes care of merging and adjusting indices in edge_index. Can you run batch.validate() to confirm your data is correctly formatted? Can you also run with CUDA_LAUNCH_BLOCKING=1? It looks like the real issue may happen earlier.

[sachinmloecher]

I tried batch.validate() but it did nothing, and running with CUDA_LAUNCH_BLOCKING=1 throws a slightly different error:

C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\cuda\ScatterGatherKernel.cu:145: block: [23,0,0], thread: [95,0,0] Assertion `idx_dim >= 0 && idx_dim < index_size && "index out of bounds"` failed.
Traceback (most recent call last):
  File "C:\Users\sachi\SensoriumDecoding\src\models\trainGCN.py", line 316, in <module>
    model = main(args)
            ^^^^^^^^^^
  File "C:\Users\sachi\SensoriumDecoding\src\models\trainGCN.py", line 280, in main
    train_result = train(
                   ^^^^^^
  File "C:\Users\sachi\SensoriumDecoding\src\models\trainGCN.py", line 169, in train
    result = train_step(
             ^^^^^^^^^^^
  File "C:\Users\sachi\SensoriumDecoding\src\models\trainGCN.py", line 129, in train_step
    y_pred = model(
             ^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\SensoriumDecoding\src\models\model.py", line 27, in forward
    mouse_output = self.mouse_mlp_dict[mouse_id](x, edge_index)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\SensoriumDecoding\src\models\mouseGCN.py", line 36, in forward
    x = gcn_layer(x, edge_index)
        ^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch_geometric\nn\conv\gcn_conv.py", line 241, in forward
    edge_index, edge_weight = gcn_norm(  # yapf: disable
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch_geometric\nn\conv\gcn_conv.py", line 108, in gcn_norm
    deg = scatter(edge_weight, idx, dim=0, dim_size=num_nodes, reduce='sum')
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\sachi\miniconda3\Lib\site-packages\torch_geometric\utils\_scatter.py", line 75, in scatter
    return src.new_zeros(size).scatter_add_(dim, index, src)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Should I be using the edge_index of one sample for the whole batch? Or let edge_index grow with batch_size?

def micro_batching(batch: Data, micro_batch_size: int):
    num_samples = len(batch.image)
    num_edges = batch.edge_index.size(1)
    edges_per_sample = num_edges // num_samples
    for i in range(0, num_samples, micro_batch_size):
        start_edge = i * edges_per_sample
        end_edge = (i + micro_batch_size) * edges_per_sample
        micro_batch = Data(
            x=batch.x[i:i+micro_batch_size],
            edge_index=batch.edge_index[:, 0:edges_per_sample],
            # edge_index=batch.edge_index[:, start_edge:end_edge],
            image=batch.image[i:i+micro_batch_size],
            behavior=batch.behavior[i:i+micro_batch_size],
            pupil_center=batch.pupil_center[i:i+micro_batch_size],
            image_id=batch.image_id[i:i+micro_batch_size],
            trial_id=batch.trial_id[i:i+micro_batch_size],
            mouse_id=batch.mouse_id[i:i+micro_batch_size]
        )
        yield micro_batch

[rusty1s]

Is this resolved or why is it closed? What's the answer to your question?

[sachinmloecher]

Yes its resolved, I did not actually figure out what was broken with my microbatching function but instead realized I could train without it. Normal batching worked perfectly fine as you mentioned with merging edge_index and x.