A model that implements multiple pooling methods in PyG

[SongZan222]

This model realizes the pooling mode of TOP-k, SAG, ASA and PAN in PyG. Write the model as a.py file that can be run with one click. The function of neural network can be controlled by the hyperparameter in the file header. At the same time, a large number of auxiliary functions are written in the file, which can help users quickly count the data and visualize the data set.

I would like to express my special thanks to @rusty1s , the author of PyG community, for his great help to me. Although he is on the other side of the earth, his help helped me, a not so good learner, successfully complete my undergraduate thesis and this model.

As the result of this period of learning, I hope this model can make some contribution to other learners and PyG community.😊

import numpy
import numpy as np
import torch
import torch_geometric
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader
from torch_geometric.graphgym import register_pooling
from torch_geometric.nn import GCNConv
from torch_geometric.nn import GIN
from torch_geometric.nn import MessagePassing
from torch_geometric.nn.pool import SAGPooling
from torch_geometric.nn.pool import PANPooling
from torch_geometric.nn.pool import ASAPooling
from torch_geometric.nn.pool import TopKPooling
from torch_geometric.nn.conv import PANConv
import torch_scatter
import scipy
from torch_geometric.nn import aggr
import matplotlib.pyplot as plt
import networkx
from timeit import default_timer as timer
from torch_geometric.nn import pool
from threading import Thread





# 0.0 Hyper-parameters 超参数
DATASETS_PATH = r'C:\Users\Administrator\Desktop\data'
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")   # Use GPU first 优先使用GPU
DATASET_NAME = 'DD'    # Dataset name in TUDataset TUDataset中的数据集名称
LEARNING_RATE = 0.01   # Learning rate 学习率
HIDDEN_DIMENSION = 16  # Hidden layer dimension 隐藏层维度
EPOCH = 64            # epoch Recommended value greater than 500 建议大于500
TEST_TIMES = 10       # Number of tests  测试次数 Suggestion 100
BATCH_SIZE_TRAIN = 64  # batch size of the training set 训练集batch大小
BATCH_SIZE_TEST = 500  # batch size of the testing set 测试集batch大小
RATIO = 0.75           # Discard rates in pooled models 池化模型中的丢弃率
DROP_OUT = 0.05
POOLING = 'ASA'        # Pooling system 'ASA'/'SAG'/'TOPk'  池化方式
MODEL_GCN = 'GIN'      # 'GIN' Convolution layer 卷积层
GLOBAL_POOLING ='g'    # 'h' or 'g'   Global pooling or hierarchical pooling
DYNAMIC_LEARNING_RATE = False   # Whether to adjust the learning rate with the epoch 是否随epoch调整学习率
SAVE_TRAINING_RESULTS = True   # Whether to store loss changes during training 是否存储训练过程中的loss变化
ACC = 0
DVA = 0
'''
def train_flag():
    global NODE_LABELS
    if DATASET_NAME == 'DD':
        return 900
    if DATASET_NAME == 'PROTEINS':
        return 900
    if DATASET_NAME == 'NCI1': # 平均节点20个左右，全局池化效果高6个点！
        return 3800
    if DATASET_NAME == 'NCI109': # 平均节点30个，全局池化效果高10个点！
        return 3800
    if DATASET_NAME == 'FRANKENSTEIN':
        return 4000
    if DATASET_NAME == 'ENZYMES':
        return 500
    if DATASET_NAME == 'PTC_FM':
        return 256
    if DATASET_NAME == 'MUTAG':
        return 160
    if DATASET_NAME == 'COLLAB':
        return 4500
    if DATASET_NAME == 'IMDB-BINARY':
        NODE_LABELS = True
        return 800
    if DATASET_NAME == 'REDDIT-BINARY':
        NODE_LABELS = True
        return 1800
    if DATASET_NAME == 'IMDB-MULTI':
        NODE_LABELS = True
        return 1300
    if DATASET_NAME == 'AIDS':
        NODE_LABELS = True
        return 1800
    if DATASET_NAME == 'REDDIT-MULTI-5K':
        NODE_LABELS = True
        return 4800
    if DATASET_NAME == 'SYNTHETIC':
        return 270
    if DATASET_NAME == 'FIRSTMM_DB':
        return 30
'''
Train_Flag = {
    'DD':900,
    'PROTEINS':900,
    'NCI1':3800,
    'NCI109':3800,
    'FRANKENSTEIN':4000,
    'ENZYMES':500,
    'PTC_FM':256,
    'MUTAG':160,
}
TRAIN_FLAG = Train_Flag[DATASET_NAME]
def get_test_item():
    return (POOLING+'_'+GLOBAL_POOLING )

TEST_ITEM = get_test_item()
# --------------------------------- #
# 0.1 Auxiliary function 辅助函数
# Draw a line graph  绘制折线图
def plt_line_graph(list_x, list_y1):
    '''

    :param list_x: The horizontal axis of the line graph，折线图的横轴
    :param list_y1: The first set of data that needs to be drawn 第一组数据
    :param list_y2: The second set of data that needs to be drawn 第二组数据
    :return: Draw the line graph
    '''
    plt.plot(
        list_x,
        list_y1,
        label="acc",
        color='#58006E'
    )
    # y1 Average value   绘制y1平均值
    print(list_x)
    print(list_y1)
    y1_average = [sum(list_y1) / (len(list_y1)), sum(list_y1) / (len(list_y1))]
    x_first_and_end = [list_x[0], list_x[-1]]
    print('the average acc is ')
    print(y1_average[0])
    plt.plot(
        x_first_and_end,
        y1_average,
        label="average acc:"+str(format(y1_average[0],'.4f')),
        color='#E5439A',
        linestyle='-.',
        markersize=15
    )
    # 绘制y1标准差
    np.std(list_y1, ddof=1)
    avge =sum(list_y1) / (len(list_y1))
    y1_std_down = [avge-np.std(list_y1, ddof=1), avge-np.std(list_y1, ddof=1)]
    y1_std_up = [avge+np.std(list_y1, ddof=1), avge+np.std(list_y1, ddof=1)]
    x_first_and_end = [list_x[0], list_x[-1]]
    print('the standard deviation is ')
    print(y1_std_down[0])
    plt.plot(
        x_first_and_end,
        y1_std_down,
        color='#FC9F81',
        label = 'standard deviation:±'+str(format(np.std(list_y1, ddof=1),'.4f')),
        linestyle='-.',
        markersize=15
    )
    plt.plot(
        x_first_and_end,
        y1_std_up,
        color='#FC9F81',
        linestyle='-.',
        markersize=15
    )
    # y1_std = []
    # )
    # plt.plot(
    #     list_x,
    #     list_y2,
    #     label="Accuracy Rate after N epochs",
    #     color='#F6573D',
    # )
    # #  y2 Average value 绘制y2平均值
    # y2_average = [sum(list_y2) / len(list_y2), sum(list_y2) / len(list_y2)]
    # x_first_and_end = [list_x[0], list_x[-1]]
    # plt.plot(
    #     x_first_and_end,
    #     y2_average,
    #     label="average accuracy",
    #     color='#FC9F81',
    #     linestyle='-.',
    #     markersize=15
    # )
    plt.legend(loc='best')
    # Control Y-axis range and step size 控制y轴范围和步长
    # plt.yticks(y_ticks[10:90:10])
    plt.grid(  # Background dotted line grid 背景虚线网格
        True,
        linestyle='--',
        alpha=0.25
    )
    # The name of the axis and the name of the chart 轴的名称和表的名称
    plt.xlabel("idx", fontdict={'size': 16})
    plt.ylabel("loss and accuracy", fontdict={'size': 16})
    plt.title(TEST_ITEM, fontdict={'size': 16})
    print(get_test_item()+'_'+DATASET_NAME+'_'+format(EPOCH))
    print(str(format(y1_average[0],'.4f'))+'±'+str(format(np.std(list_y1, ddof=1),'.4f')))
    global ACC
    global DVA
    ACC = str(format(y1_average[0],'.3f'))
    DVA = str(format(np.std(list_y1, ddof=1),'.3f'))
    # plt.show()

# Drawing matrix 绘制矩阵
def plt_matrix(m, name='the image of matrix'):
    m = plt.matshow(m, cmap=plt.cm.RdPu)
    plt.colorbar(m.colorbar, fraction=0.25)
    plt.title(name)
    plt.show()

# Plot the graph data structure 绘制图数据结构
def plt_dataset_sample(dataset, idx):
    '''

    :param dataset: The data set to display 要显示的数据集
    :param idx: Sample with serial number idx 绘制第idx个样本
    :return:Draw the adjacency matrix and the image 绘制邻接矩阵和图像
    '''
    data = dataset
    G = torch_geometric.utils.to_networkx(data[idx])
    print(type(G))
    networkx.draw_networkx(
        G=G,
        pos=networkx.spring_layout(G),
        node_size=150,
        cmap=plt.cm.RdPu,
        width=1,
        style='solid',
        label='A graph in the data set'
    )
    plt.show()

def save_list(list,name):
    name = r'training_result\_' + name
    b = np.array(list)
    np.save(name, b)

def open_list(name):
    c = np.load(name)
    return c.tolist()
# --------------------------------- #

# 1. Prepare data set  准备数据集
dataset = TUDataset(
    root=DATASETS_PATH,
    name=DATASET_NAME,
    transform=torch_geometric.transforms.ToDevice(DEVICE)
)


# prepare Datalorder
def get_dataloader(name, is_train=True):
    dataset = TUDataset(
        root=DATASETS_PATH,
        name=name,
        transform=torch_geometric.transforms.ToDevice(DEVICE)
    ).shuffle()
    train_dataset = dataset[:TRAIN_FLAG]
    test_dataset = dataset[TRAIN_FLAG:]
    loader_train = DataLoader(
        train_dataset,
        shuffle=False,
        batch_size=BATCH_SIZE_TRAIN
    )
    loader_test = DataLoader(
        test_dataset,
        shuffle=False,
        batch_size=BATCH_SIZE_TEST
    )
    if is_train:
        return loader_train
    else:
        return loader_test
# --------------------------------- #

# 2.Defined network 定义网络

# Read Out 函数
sum_aggr = aggr.SumAggregation().to(DEVICE)
avg_aggr = aggr.MeanAggregation().to(DEVICE)
max_aggr = aggr.MaxAggregation().to(DEVICE)


# Network definition

# global
# global-SAG
class Net_SAG_g(torch.nn.Module):
    def __init__(self):
        super(Net_SAG_g, self).__init__()

        self.conv1 = GIN(
            in_channels = dataset.num_node_features,
            hidden_channels = HIDDEN_DIMENSION,
            out_channels = HIDDEN_DIMENSION,
            num_layers = 3,
            dropout = DROP_OUT
        )
        self.pool = SAGPooling(HIDDEN_DIMENSION, 0.75)
        self.fc1 = torch.nn.Linear(HIDDEN_DIMENSION, 16)
        self.fc2 = torch.nn.Linear(16, 16)
        self.fc3 = torch.nn.Linear(16, dataset.num_classes)
    def forward(self, data):
        x = data.x
        edge_index = data.edge_index
        x = self.conv1(x, edge_index)
        x = torch.nn.functional.relu(x)
        x,edge_index,_,batch,_,_ = self.pool(
            x=x,
            edge_index=edge_index,
            batch=data.batch
        )
        x = sum_aggr(
            x=x,
            index= batch
        )
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return torch.nn.functional.log_softmax(x, dim=-1)
# global-TOPk
class Net_TOPk_g(torch.nn.Module):
    def __init__(self):
        super(Net_TOPk_g, self).__init__()

        self.conv1 = GIN(
            in_channels = dataset.num_node_features,
            hidden_channels = HIDDEN_DIMENSION,
            out_channels = HIDDEN_DIMENSION,
            num_layers = 3,
            dropout = DROP_OUT
        )
        self.pool = TopKPooling(in_channels=HIDDEN_DIMENSION,ratio=RATIO)
        self.fc1 = torch.nn.Linear(HIDDEN_DIMENSION, 16)
        self.fc2 = torch.nn.Linear(16, 16)
        self.fc3 = torch.nn.Linear(16, dataset.num_classes)
    def forward(self, data):
        x = data.x
        edge_index = data.edge_index
        x = self.conv1(x, edge_index)
        x = torch.nn.functional.relu(x)
        x,edge_index,_,batch,index,_ = self.pool(
            x=x,
            batch=data.batch,
            edge_index = data.edge_index,
        )
        x = sum_aggr(
            x=x,
            index= batch
        )
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return torch.nn.functional.log_softmax(x, dim=-1)
# global-SUM
class Net_SUM_g(torch.nn.Module):
    def __init__(self):
        super(Net_SUM_g, self).__init__()
        self.conv = GIN(
            in_channels  = dataset.num_node_features,
            hidden_channels = HIDDEN_DIMENSION,
            out_channels= HIDDEN_DIMENSION,
            num_layers=3,
            dropout=DROP_OUT
        )
        self.fc1 = torch.nn.Linear(HIDDEN_DIMENSION, 16)
        self.fc2 = torch.nn.Linear(16, 16)
        self.fc3 = torch.nn.Linear(16, dataset.num_classes)

    def forward(self,data):
        x= data.x
        edge_index = data.edge_index
        x = self.conv(x,edge_index)
        x = torch.nn.functional.relu(x)
        x = sum_aggr(
            x =x,
            index=data.batch
        )
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return torch.nn.functional.log_softmax(x, dim=-1)
# global-MEAN
class Net_MEAN_g(torch.nn.Module):
    def __init__(self):
        super(Net_MEAN_g, self).__init__()
        self.conv = GIN(
            in_channels  = dataset.num_node_features,
            hidden_channels = HIDDEN_DIMENSION,
            out_channels= HIDDEN_DIMENSION,
            num_layers=3,
            dropout=DROP_OUT
        )
        self.fc1 = torch.nn.Linear(HIDDEN_DIMENSION, 16)
        self.fc2 = torch.nn.Linear(16, 16)
        self.fc3 = torch.nn.Linear(16, dataset.num_classes)

    def forward(self,data):
        x= data.x
        edge_index = data.edge_index
        x = self.conv(x,edge_index)
        x = torch.nn.functional.relu(x)
        x = avg_aggr(
            x = x,
            index=data.batch
        )
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return torch.nn.functional.log_softmax(x, dim=-1)
# global-MAX
class Net_MAX_g(torch.nn.Module):
    def __init__(self):
        super(Net_MAX_g, self).__init__()
        self.conv = GIN(
            in_channels  = dataset.num_node_features,
            hidden_channels = HIDDEN_DIMENSION,
            out_channels= HIDDEN_DIMENSION,
            num_layers=3,
            dropout=DROP_OUT
        )
        self.fc1 = torch.nn.Linear(HIDDEN_DIMENSION, 16)
        self.fc2 = torch.nn.Linear(16, 16)
        self.fc3 = torch.nn.Linear(16, dataset.num_classes)

    def forward(self,data):
        x= data.x
        edge_index = data.edge_index
        x = self.conv(x,edge_index)
        x = torch.nn.functional.relu(x)
        x = max_aggr(
            x = x,
            index=data.batch
        )
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return torch.nn.functional.log_softmax(x, dim=-1)
# global-ASA
class Net_ASA_g(torch.nn.Module):
    def __init__(self):
        super(Net_ASA_g, self).__init__()
        self.conv1 = GIN(
            in_channels=dataset.num_node_features,
            hidden_channels=HIDDEN_DIMENSION,
            out_channels=HIDDEN_DIMENSION,
            num_layers=3,
            dropout=DROP_OUT
        )
        self.pool1 = ASAPooling(HIDDEN_DIMENSION)
        self.fc1 = torch.nn.Linear(HIDDEN_DIMENSION, HIDDEN_DIMENSION)
        self.fc2 = torch.nn.Linear(HIDDEN_DIMENSION, dataset.num_node_features)

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch
        x = self.conv1(x, edge_index)
        x = torch.relu(x)
        x, edge_index, _, batch, index = self.pool1(x, edge_index, None, batch)
        x = torch_geometric.nn.global_mean_pool(x,batch)
        x = self.fc1(x)
        x = self.fc2(x)
        return torch.nn.functional.log_softmax(x, dim=-1)
# global-PAN
class Net_PAN_g(torch.nn.Module):
    def __init__(self):
        super(Net_PAN_g, self).__init__()

        self.conv1 = PANConv(dataset.num_node_features, HIDDEN_DIMENSION,0)
        self.conv2 = PANConv(HIDDEN_DIMENSION, HIDDEN_DIMENSION,0)
        self.conv3 = PANConv(HIDDEN_DIMENSION, HIDDEN_DIMENSION,0)
        self.pool1 = PANPooling(HIDDEN_DIMENSION,ratio=0.5)
        self.fc1 = torch.nn.Linear(HIDDEN_DIMENSION,HIDDEN_DIMENSION)
        self.fc2 = torch.nn.Linear(HIDDEN_DIMENSION,dataset.num_classes)
    def forward(self, data):
        x = data.x
        edge_index = data.edge_index
        x, M= self.conv1(x, edge_index)
        x, M= self.conv2(x, edge_index)
        x, M = self.conv3(x, edge_index)
        x, edge_index, edge_weight, batch, perm, score = self.pool1(x,M,data.batch)
        y =avg_aggr(
            x,
            index=batch,
        )
        y = self.fc1(y)
        y = self.fc2(y)
        return torch.nn.functional.log_softmax(y, dim=-1)

# hierarchy
class SAG_GIN_h(torch.nn.Module):
    def __init__(self,c_in,c_out,ratio):
        super(SAG_GIN_h,self).__init__()
        self.pool = SAGPooling(c_in,ratio)
        self.conv = GIN(
            in_channels = c_in,
            hidden_channels = c_out,
            out_channels = c_out,
            num_layers = 1,
            dropout = DROP_OUT
        )
        self.readout = sum_aggr

    def forward(self,X,Edge,Batch):
        x, edge_index,batch = X,Edge,Batch
        x,edge_index,_,batch,_,_= self.pool(
            x=x,
            edge_index=edge_index,
            batch=batch
        )
        y = self.conv(x,edge_index)
        y = sum_aggr(
            x = y,
            index= batch,
        )
        return x,edge_index,batch,y
class SAG_TOPk_h(torch.nn.Module):
    def __init__(self,c_in,c_out,ratio):
        super(SAG_TOPk_h,self).__init__()
        self.pool = TopKPooling(c_in,ratio)
        self.conv = GIN(
            in_channels = c_in,
            hidden_channels = c_out,
            out_channels = c_out,
            num_layers = 1,
            dropout = DROP_OUT
        )
        self.readout = sum_aggr

    def forward(self,X,Edge,Batch):
        x, edge_index,batch = X,Edge,Batch
        x,edge_index,_,batch,_,_= self.pool(
            x=x,
            edge_index=edge_index,
            batch=batch
        )
        y = self.conv(x,edge_index)
        y = sum_aggr(
            x = y,
            index= batch,
        )
        return x,edge_index,batch,y
# ---------------- #
# hierarchy-SAG
class Net_SAG_h(torch.nn.Module):
    def __init__(self):
        super(Net_SAG_h, self).__init__()
        in_dimension = dataset.num_node_features
        self.SAGh_1 = SAG_GIN_h(in_dimension, HIDDEN_DIMENSION, RATIO)
        self.SAGh_2 = SAG_GIN_h(dataset.num_node_features, HIDDEN_DIMENSION, RATIO)
        self.SAGh_3 = SAG_GIN_h(dataset.num_node_features, HIDDEN_DIMENSION, RATIO)
        self.fc1 = torch.nn.Linear(3 * HIDDEN_DIMENSION, HIDDEN_DIMENSION)
        self.fc2 = torch.nn.Linear(HIDDEN_DIMENSION, dataset.num_classes)

    def forward(self, data):
        Edge_index, Batch = data.edge_index,data.batch
        X = data.x
        X, Edge_index, Batch, y_1 = self.SAGh_1(X ,Edge_index, Batch)
        X, Edge_index, Batch, y_2 = self.SAGh_2(X ,Edge_index, Batch)
        X, Edge_index, Batch, y_3 = self.SAGh_3(X, Edge_index, Batch)
        y = torch.cat((y_1,y_2,y_3),dim=1)
        y = self.fc1(y)
        y = self.fc2(y)
        return torch.nn.functional.log_softmax(y, dim=-1)
# hierarchy-TOPk
class Net_TOPk_h(torch.nn.Module):
    def __init__(self):
        super(Net_TOPk_h, self).__init__()
        in_dimension = dataset.num_node_features
        self.SAGh_1 = SAG_TOPk_h(in_dimension, HIDDEN_DIMENSION, RATIO)
        self.SAGh_2 = SAG_TOPk_h(dataset.num_node_features, HIDDEN_DIMENSION, RATIO)
        self.SAGh_3 = SAG_TOPk_h(dataset.num_node_features, HIDDEN_DIMENSION, RATIO)
        self.fc1 = torch.nn.Linear(3 * HIDDEN_DIMENSION, HIDDEN_DIMENSION)
        self.fc2 = torch.nn.Linear(HIDDEN_DIMENSION, dataset.num_classes)

    def forward(self, data):
        Edge_index, Batch = data.edge_index,data.batch
        X = data.x
        X, Edge_index, Batch, y_1 = self.SAGh_1(X ,Edge_index, Batch)
        X, Edge_index, Batch, y_2 = self.SAGh_2(X ,Edge_index, Batch)
        X, Edge_index, Batch, y_3 = self.SAGh_3(X, Edge_index, Batch)
        y = torch.cat((y_1,y_2,y_3),dim=1)
        y = self.fc1(y)
        y = self.fc2(y)
        return torch.nn.functional.log_softmax(y, dim=-1)
# hierarchy-ASA
class Net_ASA_h(torch.nn.Module):
    def __init__(self):
        super(Net_ASA_h, self).__init__()
        self.conv1 = GIN(
            in_channels=dataset.num_node_features,
            hidden_channels=HIDDEN_DIMENSION,
            out_channels=HIDDEN_DIMENSION,
            num_layers=1,
            dropout=DROP_OUT
        )
        self.pool1 = ASAPooling(HIDDEN_DIMENSION)
        self.conv2 = GIN(
            in_channels=HIDDEN_DIMENSION,
            hidden_channels=HIDDEN_DIMENSION,
            out_channels=HIDDEN_DIMENSION,
            num_layers=1,
            dropout=DROP_OUT
        )
        self.pool2 = ASAPooling(HIDDEN_DIMENSION)
        self.fc1 = torch.nn.Linear(HIDDEN_DIMENSION, HIDDEN_DIMENSION)
        self.fc2 = torch.nn.Linear(HIDDEN_DIMENSION, dataset.num_node_features)

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch

        x = self.conv1(x, edge_index)
        x, edge_index, _, batch, index = self.pool1(x, edge_index, None, batch)
        y1 = torch_geometric.nn.global_mean_pool(x,batch)

        x = self.conv2(x, edge_index)
        x, edge_index, _, batch, index = self.pool2(x, edge_index, None, batch)

        y2 = torch_geometric.nn.global_mean_pool(x,batch)
        y = torch.cat((y1,y2),dim=1)
        y = self.fc1(y)
        y = self.fc2(y)
        return torch.nn.functional.log_softmax(y, dim=-1)


# --------------------------------- #
# models
# --------------------------------- #


print(TEST_ITEM)
DIC = {
    #-------global-----------#
    'SUM_g':Net_SUM_g().to(DEVICE),
    'MEAN_g':Net_MEAN_g().to(DEVICE),
    'MAX_g':Net_MAX_g().to(DEVICE),
    'SAG_g':Net_SAG_g().to(DEVICE),
    'TOPk_g':Net_TOPk_g().to(DEVICE),
    'ASA_g':Net_ASA_g().to(DEVICE),
    'PAN_g':Net_PAN_g().to(DEVICE),

    # -------hierarchy-------#
    'SAG_h':Net_SAG_h().to(DEVICE),
    'TOPk_h':Net_TOPk_h().to(DEVICE),
    'ASA_h':Net_ASA_h().to(DEVICE),
}
model = DIC[TEST_ITEM]
print(model)
# --------------------------------- #
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
scheduler = torch.optim.lr_scheduler.MultiStepLR(
    optimizer,
    milestones=[150,200,300,400],
    gamma=0.5,
)
# --------------------------------- #

# train and test
# train
def train():
    epoch_counter = []
    loss_collector = []
    for j in range(EPOCH):
        data_loader = get_dataloader(DATASET_NAME)
        counter = 0
        for idx in data_loader:
            output = model(idx)
            optimizer.zero_grad()
            loss = torch.nn.functional.nll_loss(output, idx.y)
            loss.backward()
            optimizer.step()
            counter += 1
            if counter % 2 == 0:
                print(
                    j,
                    'loss:',
                    loss.item(),
                    'learn:',
                    format(optimizer.state_dict()['param_groups'][0]['lr']),
                )
        loss_collector.append(loss.item())
        epoch_counter.append(j)
        if DYNAMIC_LEARNING_RATE:
            scheduler.step()
    if SAVE_TRAINING_RESULTS:
        save_list(epoch_counter,str(get_test_item()+'_'+str(EPOCH))+'_'+'epoch')
        save_list(loss_collector,str(get_test_item()+'_'+str(EPOCH))+'_'+'loss')


# -------------------------------- #
# test
def test():
    model.eval()
    counter = 1
    test_counter = []
    test_acc_collecter = []
    for j in range(TEST_TIMES):
        data_loader = get_dataloader(DATASET_NAME, is_train=False)
        for batch in data_loader:
            out = model(batch).max(1)[1]
            test_acc = torch.eq(
                out, batch.y
            )
            test_counter.append(counter)
            counter += 1
            test_acc_collecter.append(test_acc.float().mean().item())
            print(batch.y)
            print(out)
            print(test_acc.float().mean())
    plt_line_graph(test_counter, test_acc_collecter)
# --------------------------------- #

# --------------------------------- #

TIC = timer()
train()
TOC = timer()
test()
TEST_DONE = timer()
train_times = format(TOC-TIC,".2f")
test_times = format(TEST_DONE-TOC,".2f")
print("training_Takes_",format(TOC-TIC,".2f"),"_S")
print("testing_Takes_",format(TEST_DONE-TOC,".2f"),"_S")
print(get_test_item()+'_'+DATASET_NAME+'_'+str(EPOCH)+'_'+ACC+'_'+DVA+'_'+train_times+'_'+test_times)
# --------------------------------- #

[rusty1s]

Wonderful, thanks for posting :)