model.py
import torch. nn as nn
from torch_geometric. nn import GATConv
import torch. nn. functional as F
class gat_cls ( nn. Module) : def __init__ ( self, in_dim, hid_dim, out_dim, dropout_size= 0.5 ) : super ( gat_cls, self) . __init__( ) self. conv1 = GATConv( in_dim, hid_dim) self. conv2 = GATConv( hid_dim, hid_dim) self. fc = nn. Linear( hid_dim, out_dim) self. relu = nn. ReLU( ) self. dropout_size = dropout_sizedef forward ( self, x, edge_index) : x = self. conv1( x, edge_index) x = F. dropout( x, p= self. dropout_size, training= self. training) x = self. relu( x) x = self. conv2( x, edge_index) x = self. relu( x) x = self. fc( x) return x
main.py
import torch
import torch. nn as nn
from torch_geometric. datasets import Planetoid
from model import gat_cls
import torch. optim as optim
dataset = Planetoid( root= './data/Cora' , name= 'Cora' )
print ( dataset[ 0 ] )
cora_data = dataset[ 0 ] epochs = 50
lr = 1e-3
weight_decay = 5e-3
momentum = 0.5
hidden_dim = 128
output_dim = 7 net = gat_cls( cora_data. x. shape[ 1 ] , hidden_dim, output_dim)
optimizer = optim. AdamW( net. parameters( ) , lr= lr, weight_decay= weight_decay)
criterion = nn. CrossEntropyLoss( )
print ( "****************Begin Training****************" )
net. train( )
for epoch in range ( epochs) : out = net( cora_data. x, cora_data. edge_index) optimizer. zero_grad( ) loss_train = criterion( out[ cora_data. train_mask] , cora_data. y[ cora_data. train_mask] ) loss_val = criterion( out[ cora_data. val_mask] , cora_data. y[ cora_data. val_mask] ) loss_train. backward( ) print ( 'epoch' , epoch+ 1 , 'loss-train {:.2f}' . format ( loss_train) , 'loss-val {:.2f}' . format ( loss_val) ) optimizer. step( ) net. eval ( )
out = net( cora_data. x, cora_data. edge_index)
loss_test = criterion( out[ cora_data. test_mask] , cora_data. y[ cora_data. test_mask] )
_, pred = torch. max ( out, dim= 1 )
pred_label = pred[ cora_data. test_mask]
true_label = cora_data. y[ cora_data. test_mask]
acc = sum ( pred_label== true_label) / len ( pred_label)
print ( "****************Begin Testing****************" )
print ( 'loss-test {:.2f}' . format ( loss_test) , 'acc {:.2f}' . format ( acc) )
参数设置
epochs = 50
lr = 1e-3
weight_decay = 5e-3
momentum = 0.5
hidden_dim = 128
output_dim = 7
运行图