Dataset to GPU

[anthonysirico]

Hello, I have been successfully running my models on the CPU and i would now like to run them on the GPU. I have created my own dataset using the following:

class CircuitDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(CircuitDataset, self).__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        return 'shuffled_data.mat'

    @property
    def processed_file_names(self):
        self.raw_data = loadmat(self.raw_paths[0], squeeze_me=True)
        self.data = self.raw_data['Graphs']
        return [f'data_{i}.pt' for i in range(len(self.data))]

    def download(self):
        pass

    def process(self):
        self.raw_data = loadmat(self.raw_paths[0], squeeze_me=True)
        self.data = pd.DataFrame(self.raw_data['Graphs'])

        for i in range(len(self.data)):
            self.data['Ln'][i] = dict(enumerate(self.data['Ln'][i]))

        idx = 0
        for index, cir in tqdm(self.data.iterrows(), total=self.data.shape[0]):
            nxg = nx.Graph(self.data['A'][index])
            nx.set_node_attributes(nxg, self.data['Ln'][index], 'component')
            pt_graph = self._get_graph_object(nxg)
            pt_graph.x = self._get_node_features(nxg, self.data['Ln'][index])
            pt_graph.y = self._get_labels(self.data['Labels'][index])
            pt_graph.performance = torch.tensor(self.data['Labels'][index], dtype=torch.float)
            torch.save(pt_graph, osp.join(self.processed_dir, f'data_{idx}.pt'))
            idx += 1

    def _get_node_features(self, nxgraph, node_labels):
        betweenness = list(nx.betweenness_centrality(nxgraph).values())
        eigenvector = list(nx.eigenvector_centrality(nxgraph, max_iter=600).values())

        mapping_dict = {'C': 0, 'G': 1, 'I': 2, 'O': 3, 'P': 4, 'R': 5}
        component_labels = []

        for value in node_labels.values():
            if value in mapping_dict:
                component_labels.append(mapping_dict[value])

        all_feats = [component_labels, betweenness, eigenvector]

        all_feats = np.asarray(all_feats).transpose()

        return torch.tensor(all_feats, dtype=torch.float)

    def _get_graph_object(self, nx_graph):
        nxg = from_networkx(nx_graph)
        return nxg

    def _get_labels(self, label):
        if label <= 0.0324:
            binary = 1
        else:
            binary = 0
        binary = np.asarray([binary])

        return torch.tensor(binary, dtype=torch.int64)

    @property
    def num_node_features(self) -> int:
        return 3

    @property
    def num_classes(self) -> int:
        return 2

    def len(self):
        data = self.raw_data['Graphs']
        return data.shape[0]

    def get(self, idx: int) -> Data:
        data = torch.load(osp.join(self.processed_dir, f'data_{idx}.pt'))
        return data

    def shuffle(self, return_perm: bool = False, ) -> Union['Dataset', Tuple['Dataset', Tensor]]:
        perm = torch.randperm(len(self))
        dataset = self.index_select(perm)
        return (dataset, perm) if return_perm is True else dataset

dataset = CircuitDataset('top25')

How can i put the dataset onto the gpu? Or is there a good way to do it within the model, shown below?

class GCN(torch.nn.Module):
    def __init__(self, hidden_channels):
        super(GCN, self).__init__()
        torch.manual_seed(12345)
        self.conv1 = GraphConv(dataset.num_node_features, hidden_channels)
        self.conv2 = GraphConv(hidden_channels, hidden_channels)
        self.conv3 = GraphConv(hidden_channels, hidden_channels)
        self.lin = Linear(hidden_channels, dataset.num_classes)

    def forward(self, x, edge_index, batch):
        # 1. Obtain node embeddings
        x = self.conv1(x, edge_index)
        x = x.relu()
        x = self.conv2(x, edge_index)
        x = x.relu()
        x = self.conv3(x, edge_index)

        # 2. Readout layer
        x = global_mean_pool(x, batch)  # [batch_size, hidden_channels]

        # 3. Apply a final classifier
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.lin(x)

        return x

model = GCN(hidden_channels=64)
print(model)

optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.CrossEntropyLoss()

def train():
    model.train()

    for data in train_loader:  # Iterate in batches over the training dataset.
        out = model(data.x, data.edge_index, data.batch)  # Perform a single forward pass.
        loss = criterion(out, data.y)  # Compute the loss.
        run['train/batch/loss'].log(loss)
        loss.backward()  # Derive gradients.
        optimizer.step()  # Update parameters based on gradients.
        optimizer.zero_grad()  # Clear gradients.

def test(loader):
    model.eval()

    correct = 0
    for data in loader:  # Iterate in batches over the training/test dataset.
        out = model(data.x, data.edge_index, data.batch)
        pred = out.argmax(dim=1)  # Use the class with highest probability.
        correct += int((pred == data.y).sum())  # Check against ground-truth labels.
    return correct / len(loader.dataset) 

[EdisonLeeeee]

It is convenient to transfer a data object between CPU and GPU, e.g.,

for data in train_loader:
    data = data.to('cuda')

Or you can do it within the model

    def forward(self, x, edge_index, batch): 
        x = x.to('cuda')
        edge_index = edge_index.to('cuda')
        batch= batch.to('cuda')

[EdisonLeeeee]

You might refer to here to implement your dataset class and make it more efficient.

[anthonysirico]

So when i create the dataset below, which by the way is typically bypassed because i have the data on my computer saved already to go in separate torch files, use InMemoryDataset?

class CircuitDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(CircuitDataset, self).__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        return 'shuffled_data.mat'

    @property
    def processed_file_names(self):
        self.raw_data = loadmat(self.raw_paths[0], squeeze_me=True)
        self.data = self.raw_data['Graphs']
        return [f'data_{i}.pt' for i in range(len(self.data))]

    def download(self):
        pass

    def process(self):
        self.raw_data = loadmat(self.raw_paths[0], squeeze_me=True)
        self.data = pd.DataFrame(self.raw_data['Graphs'])

        for i in range(len(self.data)):
            self.data['Ln'][i] = dict(enumerate(self.data['Ln'][i]))

        idx = 0
        for index, cir in tqdm(self.data.iterrows(), total=self.data.shape[0]):
            nxg = nx.Graph(self.data['A'][index])
            nx.set_node_attributes(nxg, self.data['Ln'][index], 'component')
            pt_graph = self._get_graph_object(nxg)
            pt_graph.x = self._get_node_features(nxg, self.data['Ln'][index])
            pt_graph.y = self._get_labels(self.data['Labels'][index])
            pt_graph.performance = torch.tensor(self.data['Labels'][index], dtype=torch.float)
            torch.save(pt_graph, osp.join(self.processed_dir, f'data_{idx}.pt'))
            idx += 1

    def _get_node_features(self, nxgraph, node_labels):
        betweenness = list(nx.betweenness_centrality(nxgraph).values())
        eigenvector = list(nx.eigenvector_centrality(nxgraph, max_iter=600).values())

        mapping_dict = {'C': 0, 'G': 1, 'I': 2, 'O': 3, 'P': 4, 'R': 5}
        component_labels = []

        for value in node_labels.values():
            if value in mapping_dict:
                component_labels.append(mapping_dict[value])

        all_feats = [component_labels, betweenness, eigenvector]

        all_feats = np.asarray(all_feats).transpose()

        return torch.tensor(all_feats, dtype=torch.float)

    def _get_graph_object(self, nx_graph):
        nxg = from_networkx(nx_graph)
        return nxg

    def _get_labels(self, label):
        if label <= 0.24:
            binary = 1
        else:
            binary = 0
        binary = np.asarray([binary])

        return torch.tensor(binary, dtype=torch.int64)

    @property
    def num_node_features(self) -> int:
        return 3

    @property
    def num_classes(self) -> int:
        return 2

    def len(self):
        data = self.raw_data['Graphs']
        return data.shape[0]

    def get(self, idx: int) -> Data:
        data = torch.load(osp.join(self.processed_dir, f'data_{idx}.pt'))
        return data

    def shuffle(self, return_perm: bool = False, ) -> Union['Dataset', Tuple['Dataset', Tensor]]:
        perm = torch.randperm(len(self))
        dataset = self.index_select(perm)
        return (dataset, perm) if return_perm is True else dataset

[anthonysirico]

Would my code change much? or would it just go from class CircuitDataset(Dataset):
class CircuitDataset(InMemoryDataset):?

[EdisonLeeeee]

I think class CircuitDataset(InMemoryDataset) would be better. This requires some necessary changes to your code. I think the following code should work fine:

class CircuitDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None,
                 pre_filter=None):
        super().__init__(root, transform, pre_transform, pre_filter)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return 'shuffled_data.mat'

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        pass

    def process(self):
        self.raw_data = loadmat(self.raw_paths[0], squeeze_me=True)
        self.data = pd.DataFrame(self.raw_data['Graphs'])

        for i in range(len(self.data)):
            self.data['Ln'][i] = dict(enumerate(self.data['Ln'][i]))

        data_list = []
        for index, cir in tqdm(self.data.iterrows(), total=self.data.shape[0]):
            nxg = nx.Graph(self.data['A'][index])
            nx.set_node_attributes(nxg, self.data['Ln'][index], 'component')
            pt_graph = self._get_graph_object(nxg)
            pt_graph.x = self._get_node_features(nxg, self.data['Ln'][index])
            pt_graph.y = self._get_labels(self.data['Labels'][index])
            pt_graph.performance = torch.tensor(self.data['Labels'][index], dtype=torch.float)
            data_list.append(pt_graph)

        if self.pre_filter is not None:
            data_list = [d for d in data_list if self.pre_filter(d)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(d) for d in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

    def _get_node_features(self, nxgraph, node_labels):
        betweenness = list(nx.betweenness_centrality(nxgraph).values())
        eigenvector = list(nx.eigenvector_centrality(
            nxgraph, max_iter=600).values())

        mapping_dict = {'C': 0, 'G': 1, 'I': 2, 'O': 3, 'P': 4, 'R': 5}
        component_labels = []

        for value in node_labels.values():
            if value in mapping_dict:
                component_labels.append(mapping_dict[value])

        all_feats = [component_labels, betweenness, eigenvector]

        all_feats = np.asarray(all_feats).transpose()

        return torch.tensor(all_feats, dtype=torch.float)

    def _get_graph_object(self, nx_graph):
        nxg = from_networkx(nx_graph)
        return nxg

    def _get_labels(self, label):
        if label <= 0.24:
            binary = 1
        else:
            binary = 0
        binary = np.asarray([binary])

        return torch.tensor(binary, dtype=torch.int64)

    @property
    def num_node_features(self) -> int:
        return 3

    @property
    def num_classes(self) -> int:
        return 2

[anthonysirico]

Thank you for your help!