Pytorch结合PyG实现MLP过程详解

from torch.utils.data import random_split
from torch_geometric.datasets import TUDataset
dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
train_dataset, test_dataset = random_split(dataset, [len(dataset) - 1000, 1000])

import torch.nn as nn
from torch_geometric.nn import MLP
class Net(nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers=3):
        super(Net, self).__init__()
        self.num_layers = num_layers
        self.mlp = MLP([in_channels] + [hidden_channels] * (num_layers-1) + [out_channels])
    def forward(self, x):
        return self.mlp(x)

model = Net(in_channels=dataset.num_node_features, hidden_channels=128, out_channels=dataset.num_classes, num_layers=4)

import torch.optim as optim
from torch_geometric.data import DataLoader
from tqdm import tqdm
def train(model, loader, optimizer, loss_fn):
    model.train()
    correct = 0
    total_loss = 0
    for data in tqdm(loader, desc='Training'):
        optimizer.zero_grad()
        out = model(data.x)
        pred = out.argmax(dim=1)
        loss = loss_fn(out, data.y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.num_graphs
        correct += pred.eq(data.y).sum().item()
    return total_loss / len(loader.dataset), correct / len(loader.dataset)

def test(model, loader, loss_fn):
    model.eval()
    correct = 0
    total_loss = 0
    with torch.no_grad():
        for data in tqdm(loader, desc='Testing'):
            out = model(data.x)
            pred = out.argmax(dim=1)
            loss = loss_fn(out, data.y)
            total_loss += loss.item() * data.num_graphs
            correct += pred.eq(data.y).sum().item()
    return total_loss / len(loader.dataset), correct / len(loader.dataset)

if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = Net(in_channels=dataset.num_node_features, hidden_channels=128, out_channels=dataset.num_classes, num_layers=4).to(device)
    loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=128)
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    loss_fn = nn.CrossEntropyLoss()
    for epoch in range(1, 201):
        train_loss, train_acc = train(model, loader, optimizer, loss_fn)
        test_loss, test_acc = test(model, test_loader, loss_fn)
        print(f'Epoch {epoch:03D}, Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, '
              f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}')