怎么用代码求GIN模型的可训练参数量

在计算一个神经网络模型，如生成式对抗网络 (GAN) 中 GIN (Graph Isomorphism Network) 的可训练参数量时，通常需要遍历模型的所有层，并统计每层的参数。假设你有一个简单的 Python 示例，使用 PyTorch 实现了一个 GIN 模型，你可以这样做：

import torch

class GINConv(torch.nn.Module):
    # GIN 层示例
    def __init__(self, in_features, out_features):
        super(GINConv, self).__init__()
        self.linear = torch.nn.Linear(in_features, out_features)
        self.bn = torch.nn.BatchNorm1d(out_features)
        self.mlp = torch.nn.Sequential(
            torch.nn.Linear(1, 1),
            torch.nn.ReLU(),
            torch.nn.Linear(1, out_features - in_features)
        )

    def forward(self, x, edge_index, size=None):
        h = self.linear(x).mean(dim=1)  # 对每个节点应用平均池化
        h_prime = self.bn(h)
        m = self.mlp((h * h_prime).sum(dim=1))  # 计算特征融合
        return x + m.unsqueeze(-1).repeat(1, size[1], 1)

# 假设你有一个名为 model 的 GIN 模型实例
gin_model = YourGINModel()  # 定义并实例化 GIN 模型

# 获取所有层的列表
all_layers = list(gin_model.modules())

# 初始化参数计数器
total_params = 0

for layer in all_layers:
    if isinstance(layer, torch.nn.Linear):  # 如果是线性层
        total_params += layer.weight.numel() + layer.bias.numel()
    elif isinstance(layer, torch.nn.BatchNorm1d):  # 如果是 BatchNorm 层
        total_params += layer.weight.numel() + layer.bias.numel()

print(f"模型的总可训练参数量为: {total_params}")