How can I change the shape of my train_mask?

anto: <p>I’m performing node classification using the WikipediaNetwork dataset. The problem is that the shape of the train_mask is [2277,10] but I want a shape equal to [2277].<br>
The code that I’m using is the following one:</p>
<pre><code class="lang-auto">squirrel_dataset = WikipediaNetwork(root='data/WikipediaNetwork', name='squirrel', transform=None)
squirrel_data = squirrel_dataset[0]
</code></pre>
<pre><code class="lang-auto">def train():
      model.train()
      data = model.data
      optimizer.zero_grad()  # Clear gradients.
      out = model(data.x, data.edge_index)  # Perform a single forward pass.
      loss = criterion(out[data.train_mask], data.y[data.train_mask])  # Compute the loss solely based on the training nodes.
      loss.backward()  # Derive gradients.
      optimizer.step()  # Update parameters based on gradients.
      return loss
</code></pre>
<p>The error that I obtain is the following one:</p>
<pre><code class="lang-auto">The shape of the mask [2277, 10] at index 1 does not match the shape of the indexed tensor [2277, 5] at index 1
</code></pre>
<p>How can I solve this issue?</p>

J_Johnson: <p>Can you share the definition for <code>train_mask</code>?</p>

anto: <p>The train mask Is automatically generated when I call these lines and is part of the Data Object generated by the class WikipediaNetwork.</p>
<pre><code class="lang-auto">squirrel_dataset = WikipediaNetwork(root='data/WikipediaNetwork', name='squirrel', transform=None)
squirrel_data = squirrel_dataset[0]
</code></pre>

J_Johnson: <p>Suppose we have some tensor of arbitrary number of dimensions for a dummy dataset:</p>
<pre><code class="lang-auto">data_size=30000
A=torch.rand((data_size, 200, 5, 20))
</code></pre>
<p>Now suppose we create a mask vector of our data_size:</p>
<pre><code class="lang-auto">import numpy as np


b=np.arange(data_size)
</code></pre>
<p>Randomize it with:</p>
<pre><code class="lang-auto">rng = np.random.default_rng()
rng.shuffle(b)
</code></pre>
<p>And split it into train and test vectors with a split of 85% to 15%:</p>
<pre><code class="lang-auto">a_train=b[:round(0.85*np.shape(b)[0])]
a_test=b[round(0.85*np.shape(b)[0]):]
</code></pre>
<p>You could then get your training set by:</p>
<pre><code class="lang-auto">train_data=A[a_train,:,:,:]
</code></pre>
<p>In other words, for some reason your mask is not only on the first dim, which is the problem. Maybe they used an older version of Pytorch when developing their scripts.</p>
<p>Without seeing their definition for this object call, it’s not possible to identify the error.</p>

anto: <p>This is the source code that I have: <a href="https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/datasets/webkb.html#WebKB" class="inline-onebox" rel="noopener nofollow ugc">torch_geometric.datasets.webkb — pytorch_geometric documentation</a><br>
I hope can be useful to solve the issue.</p>

J_Johnson: <p>Your code seems to have loaded the dataset, and specified what data you want. But I do not see where you have put your data into a dataloader. Please see here:</p>
<p><a href="https://pytorch-geometric.readthedocs.io/en/latest/modules/loader.html#torch_geometric.loader.DataLoader" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch-geometric.readthedocs.io/en/latest/modules/loader.html#torch_geometric.loader.DataLoader</a></p>
<p>Copied from their example:</p>
<pre><code class="lang-auto">from torch_geometric.datasets import Planetoid
from torch_geometric.loader import NeighborLoader

data = Planetoid(path, name='Cora')[0]

loader = NeighborLoader(
    data,
    # Sample 30 neighbors for each node for 2 iterations
    num_neighbors=[30] * 2,
    # Use a batch size of 128 for sampling training nodes
    batch_size=128,
    input_nodes=data.train_mask,
)

sampled_data = next(iter(loader))
print(sampled_data.batch_size)
&gt;&gt;&gt; 128
</code></pre>

Burouj_Armgaan: <p>Hi Antonio!<img src="https://discuss.pytorch.org/images/emoji/apple/wave/3.png?v=12" title=":wave:t3:" class="emoji" alt=":wave:t3:" loading="lazy" width="20" height="20"><br>
I ran into the same problem with the <a href="https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html#torch_geometric.datasets.WebKB" rel="noopener nofollow ugc">WebKB</a> dataset.</p>
<ul>
<li>So, what’s happening here is that, instead of just one split, the authors of the dataset have made ten random splits of the dataset for train, validation and test nodes. If you print out the <code>data.train_mask</code>, you’ll see that it’s a boolean tensor. If you take <code>data.train_mask[:, 0]</code>, you’ll get data as per the 0th random split for train, val, and test nodes. It’s just for randomization.</li>
<li>So, you can do <code>data.train_mask = data.train_mask[:, 0]</code>. Similarly for val and test masks.</li>
<li>Or, you can iterate over all the splits: <code>data.train_mask[:, 0], data.train_mask[:, 1], ...</code> and so on. You’ll then average the metrics you receive over each split.</li>
</ul>
<p>Hope this helps🙂</p>

anto: <p>This solves my problem, thank you! <a class="mention" href="/u/burouj_armgaan">@Burouj_Armgaan</a></p>