AttributeError in torch_geometric.transforms


I have a problem that I cannot understand: even though a module ‘torch_geometric.transforms’ has an attribute ‘AddTrainValTestMask’ according to documentation, I cannot import it. I keep receiving an error AttributeError: module 'torch_geometric.transforms' has no attribute 'AddTrainValTestMask

My Pytorch version is 1.7.1

I took the code from here

Minimum reproducible example:

import os.path as osp

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T
from torch_geometric.nn import SplineConv

dataset = 'Cora'
transform = T.Compose([
    T.AddTrainValTestMask('train_rest', num_val=500, num_test=500),
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', dataset)
dataset = Planetoid(path, dataset, transform=transform)
data = dataset[0]

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = SplineConv(dataset.num_features, 16, dim=1, kernel_size=2)
        self.conv2 = SplineConv(16, dataset.num_classes, dim=1, kernel_size=2)

    def forward(self):
        x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
        x = F.dropout(x,
        x = F.elu(self.conv1(x, edge_index, edge_attr))
        x = F.dropout(x,
        x = self.conv2(x, edge_index, edge_attr)
        return F.log_softmax(x, dim=1)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model, data = Net().to(device),
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-3)

def train():
    F.nll_loss(model()[data.train_mask], data.y[data.train_mask]).backward()

def test():
    log_probs, accs = model(), []
    for _, mask in data('train_mask', 'test_mask'):
        pred = log_probs[mask].max(1)[1]
        acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()
    return accs

for epoch in range(1, 201):
    log = 'Epoch: {:03d}, Train: {:.4f}, Test: {:.4f}'
    print(log.format(epoch, *test()))

Can somebody explain where is the problem?

Maybe your torch_geometric version is too old. Note that this library is not shipped in the PyTorch installation, but is installed separately, so you would have to check its version.

Eventually it was the reason! Thank you very much!