【GGNN源码解析】逐步PyTorch Implementation of Gated Graph Neural Network

【GGNN源码解析】逐步PyTorch Implementation of Gated Graph Neural Network

main.py

导入必要的Python库和模块

• argparse：用于解析命令行参数。
• random：用于生成随机数。
• torch：PyTorch深度学习框架。
• nn：PyTorch的神经网络模块。
• optim：PyTorch的优化器模块。
• 从自定义模块和库中导入其他函数和类。

import argparse
import random

import torch
import torch.nn as nn
import torch.optim as optim

from model import GGNN
from utils.train import train
from utils.test import test
from utils.data.dataset import bAbIDataset
from utils.data.dataloader import bAbIDataloader

parser = argparse.ArgumentParser() # 创建了一个参数解析器，用于从命令行获取参数

# 定义了一系列命令行参数，例如task_id（bAbI任务ID）、question_id（问题类型）、workers（数据加载器工作线程数）等。这些参数将用于配置训练和测试过程。
parser.add_argument('--task_id', type=int, default=4, help='bAbI task id')
parser.add_argument('--question_id', type=int, default=0, help='question types')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=2)
parser.add_argument('--batchSize', type=int, default=10, help='input batch size')
parser.add_argument('--state_dim', type=int, default=4, help='GGNN hidden state size')
parser.add_argument('--n_steps', type=int, default=5, help='propogation steps number of GGNN')
parser.add_argument('--niter', type=int, default=10, help='number of epochs to train for')
parser.add_argument('--lr', type=float, default=0.01, help='learning rate')
parser.add_argument('--cuda', action='store_true',default="cpu", help='enables cuda')
parser.add_argument('--verbal', action='store_true', help='print training info or not')
parser.add_argument('--manualSeed', type=int, help='manual seed')

opt = parser.parse_args() # 解析命令行参数，并将解析结果存储在opt变量中。
print(opt) # 使用print(opt)打印出解析后的命令行参数，以便查看配置

 # 随机生成一个随机种子，并设置为manualSeed。
if opt.manualSeed is None:
    opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
# 使用生成的随机种子设置Python的随机数生成器和PyTorch的随机种子
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)

# 根据task_id参数构建数据集路径，这里使用的是bAbI数据集
opt.dataroot = 'babi_data/processed_1/train/%d_graphs.txt' % opt.task_id

if opt.cuda:
  # 如果opt.cuda为True，表示启用CUDA，将PyTorch模型和数据移动到GPU上
    torch.cuda.manual_seed_all(opt.manualSeed)

# main函数是整个脚本的主要逻辑入口
def main(opt):
  # 用bAbIDataset类和bAbIDataloader类创建训练集和测试集的数据加载器，这些加载器将用于加载图数据。
    train_dataset = bAbIDataset(opt.dataroot, opt.question_id, True)
    train_dataloader = bAbIDataloader(train_dataset, batch_size=opt.batchSize, \
                                      shuffle=True, num_workers=2)

    test_dataset = bAbIDataset(opt.dataroot, opt.question_id, False)
    test_dataloader = bAbIDataloader(test_dataset, batch_size=opt.batchSize, \
                                     shuffle=False, num_workers=2)
    # 设置了一些与数据集相关的模型参数，例如注释维度、边类型数量和节点数量
    opt.annotation_dim = 1  # for bAbI
    opt.n_edge_types = train_dataset.n_edge_types
    opt.n_node = train_dataset.n_node

    # 创建了一个GGNN模型的实例，并将其设置为双精度浮点数数据类型
    net = GGNN(opt)
    net.double()
    print(net)
		
    # 创建了一个GGNN模型的实例，并将其设置为双精度浮点数数据类型。
    criterion = nn.CrossEntropyLoss()
		# 如果启用了CUDA，将模型和损失函数移动到GPU上
    if opt.cuda:
        net.cuda()
        criterion.cuda()

    optimizer = optim.Adam(net.parameters(), lr=opt.lr)
		# 使用for循环进行多个训练和测试的迭代，迭代次数由opt.niter指定
    for epoch in range(0, opt.niter):
        # 在每个训练和测试迭代中，调用train函数和test函数，分别执行训练和测试过程
        train(epoch, train_dataloader, net, criterion, optimizer, opt)
        test(test_dataloader, net, criterion, optimizer, opt)


if __name__ == "__main__":
    main(opt) # 调用main(opt)函数，启动训练和测试流程。

utils

data

dataset.py

import numpy as np

def load_graphs_from_file(file_name):
    """函数用于从文件中加载图数据"""
    data_list = []
    edge_list = []
    target_list = []
    with open(file_name,'r') as f: # 打开指定文件（file_name），逐行读取文件内容
        for line in f:
            if len(line.strip()) == 0: # 通过空行将不同图的数据分隔开，将每个图的边列表和目标列表存储在data_list中
                data_list.append([edge_list,target_list]) 
                edge_list = [] # 边列表由数字列表组成，表示图中的边
                target_list = [] # 目标列表包含数字列表，表示图中的目标节点
            else:
                digits = []
                line_tokens = line.split(" ")
                if line_tokens[0] == "?":
                    for i in range(1, len(line_tokens)):
                        digits.append(int(line_tokens[i]))
                    target_list.append(digits)
                else:
                    for i in range(len(line_tokens)):
                        digits.append(int(line_tokens[i]))
                    edge_list.append(digits)
    return data_list # 函数返回包含多个图数据的列表data_list

def find_max_edge_id(data_list):
  	# 该函数用于查找数据集中的最大边ID。
    max_edge_id = 0
    # 遍历data_list中的每个图，查找每个图的边列表，找到最大的边ID并返回
    for data in data_list:
        edges = data[0]
        for item in edges:
            if item[1] > max_edge_id:
                max_edge_id = item[1]
    return max_edge_id

def find_max_node_id(data_list):
  	# 函数用于查找数据集中的最大节点ID
    max_node_id = 0
    for data in data_list:
      	# 函数用于查找数据集中的最大节点ID
        edges = data[0]
        for item in edges:
          	# 还考虑了目标节点的ID
            if item[0] > max_node_id:
                max_node_id = item[0]
            if item[2] > max_node_id:
                max_node_id = item[2]
    return max_node_id

def find_max_task_id(data_list):
  	# 还考虑了目标节点的ID
    max_node_id = 0
    for data in data_list:
      	# 遍历data_list中的每个图，查找每个图的目标列表，找到最大的任务ID并返回
        targe = data[1]
        for item in targe:
            if item[0] > max_node_id:
                max_node_id = item[0]
    return max_node_id

def split_set(data_list):
  	"""该函数用于将数据集划分为训练集和验证集"""
    n_examples = len(data_list)
    # 首先计算数据集的总示例数，并创建索引范围（idx）
    idx = range(n_examples)
    # 从索引范围中选择前50个示例作为训练集，选择后50个示例作为验证集
    train = idx[:50]
    val = idx[-50:]
    # 返回训练集和验证集的数组形式
    return np.array(data_list)[train],np.array(data_list)[val]

def data_convert(data_list, n_annotation_dim):
  	"""函数用于将数据转换为适合GGNN模型的格式"""
    # 输入数据列表data_list和注释维度n_annotation_dim
    n_nodes = find_max_node_id(data_list)
    n_tasks = find_max_task_id(data_list)
    task_data_list = []
    # 根据最大节点ID和最大任务ID创建任务数据列表task_data_list
    for i in range(n_tasks):
        task_data_list.append([])
    # 根据最大节点ID和最大任务ID创建任务数据列表task_data_list
    for item in data_list:
        edge_list = item[0]
        target_list = item[1]
        for target in target_list:
            task_type = target[0]
            task_output = target[-1]
            annotation = np.zeros([n_nodes, n_annotation_dim])
            annotation[target[1]-1][0] = 1
            task_data_list[task_type-1].append([edge_list, annotation, task_output])
    return task_data_list # 最终返回任务数据列表task_data_list

def create_adjacency_matrix(edges, n_nodes, n_edge_types):
  	"""函数用于创建邻接矩阵
    输入参数包括边列表edges、节点数n_nodes和边类型数n_edge_types
    """
    a = np.zeros([n_nodes, n_nodes * n_edge_types * 2])
    # 创建一个全零的邻接矩阵a，形状为[n_nodes, n_nodes * n_edge_types * 2]
    # 遍历边列表，为邻接矩阵中的相应位置设置值为1的条目
    for edge in edges:
        src_idx = edge[0]
        e_type = edge[1]
        tgt_idx = edge[2]
        a[tgt_idx-1][(e_type - 1) * n_nodes + src_idx - 1] =  1
        a[src_idx-1][(e_type - 1 + n_edge_types) * n_nodes + tgt_idx - 1] =  1
    return a # 返回创建的邻接矩阵a


class bAbIDataset():
    """
    Load bAbI tasks for GGNN
    该类用于加载bAbI任务的数据集
    """
    def __init__(self, path, task_id, is_train):
      	# 构造函数初始化数据集相关的信息，包括边类型数、任务数和节点数
        all_data = load_graphs_from_file(path)
        self.n_edge_types =  find_max_edge_id(all_data)
        self.n_tasks = find_max_task_id(all_data)
        self.n_node = find_max_node_id(all_data)

        all_task_train_data, all_task_val_data = split_set(all_data)
				
        # 根据训练标志is_train选择加载训练集或验证集的数据
        if is_train:
            all_task_train_data = data_convert(all_task_train_data, 1)
            self.data = all_task_train_data[task_id]
        else:
            all_task_val_data = data_convert(all_task_val_data, 1)
            self.data = all_task_val_data[task_id] # 数据以合适的格式存储在self.data中

    def __getitem__(self, index):
      	"""该方法用于从数据集中获取单个示例"""
        am = create_adjacency_matrix(self.data[index][0], self.n_node, self.n_edge_types)
        annotation = self.data[index][1]
        target = self.data[index][2] - 1
        return am, annotation, target # 返回邻接矩阵、注释和目标

    def __len__(self):
        return len(self.data) # 该方法用于获取数据集的长度，即示例的数量

dataloader.py

from torch.utils.data import DataLoader
"""导入PyTorch中的DataLoader类，这是一个用于加载和批处理数据的实用工具"""
class bAbIDataloader(DataLoader):
  """定义一个新的类bAbIDataloader，该类继承自DataLoader类，这意味着它将具有DataLoader类的所有功能和方法，并且可以通过继承来自定义或扩展其行为。"""

    def __init__(self, *args, **kwargs):
      """定义bAbIDataloader类的构造函数（__init__方法），用于初始化数据加载器对象。
*args和**kwargs是Python中的特殊参数，它们允许函数接受任意数量的位置参数（*args）和关键字参数（**kwargs）。
在这种情况下，bAbIDataloader的构造函数接受与DataLoader类相同的参数，以便将这些参数传递给DataLoader的构造函数。"""
        super(bAbIDataloader, self).__init__(*args, **kwargs)
        """使用super()函数调用父类DataLoader的构造函数，以确保bAbIDataloader类继承并正确初始化DataLoader类的功能。
*args和**kwargs参数将传递给父类的构造函数，以便初始化数据加载器"""

train.py

• epoch：当前训练的轮数。
• dataloader：用于加载训练数据的数据加载器。
• net：GGNN模型。
• criterion：损失函数，用于计算训练损失。
• optimizer：优化器，用于更新模型参数。
• opt：包含训练配置参数的对象。

---

• padding：创建一个与annotation相同大小的零张量，用于填充输入以匹配所需的状态维度。
• init_input：将annotation和padding按维度拼接以创建初始输入

import torch
from torch.autograd import Variable
"""导入PyTorch深度学习框架以及变量（Variable）模块，用于创建和处理神经网络中的变量。"""
def train(epoch, dataloader, net, criterion, optimizer, opt):
  	"""定义了一个名为train的函数，用于训练GGNN模型"""
    # 将模型设置为训练模式，这会启用批量归一化和丢弃等训练相关的操作
    net.train()
    for i, (adj_matrix, annotation, target) in enumerate(dataloader, 0):
      	# 迭代训练数据加载器，其中adj_matrix表示邻接矩阵，annotation表示节点注释，target表示目标标签
        # 清零模型的梯度，以便计算新的梯度
        net.zero_grad()

        padding = torch.zeros(len(annotation), opt.n_node, opt.state_dim - opt.annotation_dim).double()
        init_input = torch.cat((annotation, padding), 2)
        
        # 如果使用GPU（通过opt.cuda设置为True），则将数据移动到GPU上
        if opt.cuda:
            init_input = init_input.cuda()
            adj_matrix = adj_matrix.cuda()
            annotation = annotation.cuda()
            target = target.cuda()
				# init_input、adj_matrix、annotation 和 target 被包装成PyTorch的Variable对象，以便计算梯度和自动微分
        init_input = Variable(init_input)
        adj_matrix = Variable(adj_matrix)
        annotation = Variable(annotation)
        target = Variable(target)
				# 使用GGNN模型进行前向传播，计算预测输出
        output = net(init_input, annotation, adj_matrix)
				# 使用指定的损失函数criterion计算模型的预测输出与真实目标之间的损失
        loss = criterion(output, target) 

        loss.backward() # 反向传播损失，计算参数的梯度
        optimizer.step()# 使用优化器更新模型的参数，以减小损失
				# 如果opt.verbal为True，则每训练一定数量的批次会打印当前的训练进度，包括当前轮次、总轮次、当前批次和总批次以及损失值
        if i % int(len(dataloader) / 10 + 1) == 0 and opt.verbal:
            print('[%d/%d][%d/%d] Loss: %.4f' % (epoch, opt.niter, i, len(dataloader), loss.data[0]))

test.py

import torch
from torch.autograd import Variable

def test(dataloader, net, criterion, optimizer, opt):
    test_loss = 0
    correct = 0
    net.eval()  # 将模型设置为评估模式，以便禁用训练时的一些操作
    for i, (adj_matrix, annotation, target) in enumerate(dataloader, 0):
        # 创建零填充的 padding 张量，以使输入与所需的状态维度匹配
        padding = torch.zeros(len(annotation), opt.n_node, opt.state_dim - opt.annotation_dim).double()
        # 将注释和 padding 拼接在一起，创建初始化输入
        init_input = torch.cat((annotation, padding), 2)
        if opt.cuda:
            init_input = init_input.cuda()  # 如果使用 GPU，将数据移动到 GPU 上
            adj_matrix = adj_matrix.cuda()
            annotation = annotation.cuda()
            target = target.cuda()

        init_input = Variable(init_input)  # 创建 PyTorch 变量
        adj_matrix = Variable(adj_matrix)
        annotation = Variable(annotation)
        target = Variable(target)

        output = net(init_input, annotation, adj_matrix)  # 执行前向传播，计算预测输出

        # 计算测试损失
        test_loss += criterion(output, target).data[0]

        # 预测最大值的类别
        pred = output.data.max(1, keepdim=True)[1]

        # 计算正确分类的数量
        correct += pred.eq(target.data.view_as(pred)).cpu().sum()

    test_loss /= len(dataloader.dataset)  # 计算平均测试损失
    print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
        test_loss, correct, len(dataloader.dataset),
        100. * correct / len(dataloader.dataset)))

model.py

import torch
import torch.nn as nn
# PyTorch深度学习框架以及神经网络模块，用于创建神经网络

class AttrProxy(object):
    """
    Translates index lookups into attribute lookups.
    To implement some trick which able to use list of nn.Module in a nn.Module
    see https://discuss.pytorch.org/t/list-of-nn-module-in-a-nn-module/219/2
    定义了一个名为AttrProxy的类，用于将索引查找转换为属性查找。这是一种将列表中的nn.Module对象用于nn.Module中的技巧。
    """
    def __init__(self, module, prefix):
      	"""定义了一个名为AttrProxy的类，用于将索引查找转换为属性查找。这是一种将列表中的nn.Module对象用于nn.Module中的技巧。"""
        self.module = module
        self.prefix = prefix

    def __getitem__(self, i):
        return getattr(self.module, self.prefix + str(i))


class Propogator(nn.Module):
    """
    Gated Propogator for GGNN
    Using LSTM gating mechanism
    定义了名为Propogator的GGNN中的门控传播器，使用了LSTM门控机制
    """
    def __init__(self, state_dim, n_node, n_edge_types):
      	# 初始化门控传播器。接受state_dim（隐藏状态的维度）、n_node（节点数量）、n_edge_types（边类型数量）等参数
        super(Propogator, self).__init__()

        self.n_node = n_node
        self.n_edge_types = n_edge_types
				# reset_gate、update_gate 和 tansform 分别表示重置门、更新门和传播转换，都是由线性层（nn.Linear）和激活函数构成
        self.reset_gate = nn.Sequential(
            nn.Linear(state_dim*3, state_dim),
            nn.Sigmoid()
        )
        self.update_gate = nn.Sequential(
            nn.Linear(state_dim*3, state_dim),
            nn.Sigmoid()
        )
        self.tansform = nn.Sequential(
            nn.Linear(state_dim*3, state_dim),
            nn.Tanh()
        )

    def forward(self, state_in, state_out, state_cur, A):
      	"""定义了门控传播器的前向传播逻辑。接受输入state_in（传入的状态）、state_out（传出的状态）、state_cur（当前状态）、A（邻接矩阵）"""
        # 通过输入状态和邻接矩阵计算重置门、更新门和传播转换
        A_in = A[:, :, :self.n_node*self.n_edge_types]
        A_out = A[:, :, self.n_node*self.n_edge_types:]

        a_in = torch.bmm(A_in, state_in)
        a_out = torch.bmm(A_out, state_out)
        a = torch.cat((a_in, a_out, state_cur), 2)

        r = self.reset_gate(a)
        z = self.update_gate(a)
        joined_input = torch.cat((a_in, a_out, r * state_cur), 2)
        h_hat = self.tansform(joined_input)
				# 最终计算输出状态，使用门控机制来控制状态的更新
        output = (1 - z) * state_cur + z * h_hat

        return output


class GGNN(nn.Module):
    """
    Gated Graph Sequence Neural Networks (GGNN)
    Mode: SelectNode
    Implementation based on https://arxiv.org/abs/1511.05493
    定义了Gated Graph Sequence Neural Networks (GGNN) 模型
    """
    def __init__(self, opt):
      	# 初始化GGNN模型。接受包含一系列模型参数的opt参数，如状态维度、注释维度、边类型数量等
        super(GGNN, self).__init__()

        assert (opt.state_dim >= opt.annotation_dim,  \
                'state_dim must be no less than annotation_dim')
			
        self.state_dim = opt.state_dim
        self.annotation_dim = opt.annotation_dim
        self.n_edge_types = opt.n_edge_types
        self.n_node = opt.n_node
        self.n_steps = opt.n_steps

        for i in range(self.n_edge_types):
            # incoming and outgoing edge embedding
            in_fc = nn.Linear(self.state_dim, self.state_dim)
            out_fc = nn.Linear(self.state_dim, self.state_dim)
            self.add_module("in_{}".format(i), in_fc)
            self.add_module("out_{}".format(i), out_fc)

        self.in_fcs = AttrProxy(self, "in_")
        self.out_fcs = AttrProxy(self, "out_")

        # Propogation Model
        self.propogator = Propogator(self.state_dim, self.n_node, self.n_edge_types)

        # Output Model
        self.out = nn.Sequential(
            nn.Linear(self.state_dim + self.annotation_dim, self.state_dim),
            nn.Tanh(),
            nn.Linear(self.state_dim, 1)
        )
        	
        # 使用AttrProxy类将输入线性层和输出线性层按边类型进行索引，以便在多种边类型之间共享权重

        self._initialization()

    def _initialization(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                m.weight.data.normal_(0.0, 0.02)
                m.bias.data.fill_(0)

    def forward(self, prop_state, annotation, A):
      	"""定义了GGNN模型的前向传播逻辑。接受prop_state（传播状态）、annotation（注释信息）、A（邻接矩阵）"""
        for i_step in range(self.n_steps):
          	# 通过多次门控传播器的迭代，更新传播状态
            in_states = []
            out_states = []
            for i in range(self.n_edge_types):
                in_states.append(self.in_fcs[i](prop_state))
                out_states.append(self.out_fcs[i](prop_state))
            in_states = torch.stack(in_states).transpose(0, 1).contiguous()
            in_states = in_states.view(-1, self.n_node*self.n_edge_types, self.state_dim)
            out_states = torch.stack(out_states).transpose(0, 1).contiguous()
            out_states = out_states.view(-1, self.n_node*self.n_edge_types, self.state_dim)
						
            # 最后将传播状态和注释信息连接起来，并通过输出模型生成最终的输出
            prop_state = self.propogator(in_states, out_states, prop_state, A)

        join_state = torch.cat((prop_state, annotation), 2)

        output = self.out(join_state)
        output = output.sum(2)

        return output