import argparse
import sys
import torch
import time
import scipy.io as sio
import numpy as np
from torch.utils.data import TensorDataset, DataLoader
def readfile(path):
print('reading file ...')
data = sio.loadmat(path)
x_train = []
x_label = []
val_data = []
val_label = []
x_train = data['train_data']
x_label = data['train_label']
val_data = data['test_data']
val_label = data['test_label']
x_train = np.array(x_train, dtype=float)
val_data = np.array(val_data, dtype=float)
x_label = np.array(x_label, dtype=int)
val_label = np.array(val_label, dtype=int)
x_train = torch.FloatTensor(x_train)
val_data = torch.FloatTensor(val_data)
x_label = torch.LongTensor(x_label)
val_label = torch.LongTensor(val_label)
return x_train, x_label, val_data, val_label
class CNNnet(torch.nn.Module):
def __init__(self, node_number, batch_size, k_hop):
super(CNNnet,self).__init__()
self.node_number = node_number
self.batch_size = batch_size
self.k_hop = k_hop
self.aggregate_weight = torch.nn.Parameter(torch.rand(1, 1, node_number))
self.conv1 = torch.nn.Sequential(
torch.nn.Conv1d(in_channels=1,
out_channels=8,
kernel_size=3,
stride=1,
padding=1),
torch.nn.BatchNorm1d(8),
torch.nn.ReLU(),
torch.nn.MaxPool1d(kernel_size=2),
#torch.nn.AvgPool1d(kernel_size=2),
torch.nn.Dropout(0.2),
)
self.conv2 = torch.nn.Sequential(
torch.nn.Conv1d(8,16,3,1,1),
torch.nn.BatchNorm1d(16),
torch.nn.ReLU(),
torch.nn.MaxPool1d(kernel_size=2),
#torch.nn.AvgPool1d(kernel_size=2),
torch.nn.Dropout(0.2),
)
self.mlp1 = torch.nn.Sequential(
torch.nn.Linear(64*16,50),
torch.nn.Dropout(0.5),
)
self.mlp2 = torch.nn.Linear(50,2)
def forward(self, x):
tmp_x = x
for _ in range(self.k_hop):
tmp_x = torch.matmul(tmp_x, x)
x = torch.matmul(self.aggregate_weight, tmp_x)
x = self.conv1(x)
x = self.conv2(x)
x = self.mlp1(x.view(x.size(0),-1))
x = self.mlp2(x)
return x
def main():
parser = argparse.ArgumentParser(description='PyTorch graph convolutional neural net for whole-graph classification')
parser.add_argument('--dataset', type=str, default="dataset/AEF_V_0.mat", help='path of the dataset (default: data/data.mat)')
parser.add_argument('--node_number', type=int, default=256, help='node number of graph (default: 256)')
parser.add_argument('--batch_size', type=int, default=32, help='number of input size (default: 128)')
parser.add_argument('--k_hop', type=int, default=4, help='times of aggregate (default: 1)')
args = parser.parse_args()
x_train, x_label, val_data, val_label = readfile(args.dataset) # 'train.csv'
x_train = x_train.permute(2, 0, 1)
x_label = torch.squeeze(x_label, dim=1).long()
val_data = val_data.permute(2, 0, 1)
val_label = torch.squeeze(val_label, dim=1).long()
train_set = TensorDataset(x_train, x_label)
val_set = TensorDataset(val_data, val_label)
#batch_size = 128
train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=0)
val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=True, num_workers=0)
model = CNNnet(args.node_number, args.batch_size, args.k_hop)
#print(model)
model
loss = torch.nn.CrossEntropyLoss()
#para = list(model.parameters())
#print(para)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # optimize all cnn parameters
loss_func = torch.nn.CrossEntropyLoss()
best_acc = 0.0
num_epoch = 100
for epoch in range(num_epoch):
epoch_start_time = time.time()
train_acc = 0.0
train_loss = 0.0
val_acc = 0.0
val_loss = 0.0
model.train()
for i, data in enumerate(train_loader):
optimizer.zero_grad()
train_pred = model(data[0])
#print(train_pred.size())
#print(data[1].size())
batch_loss = loss(train_pred, data[1])
batch_loss.backward()
optimizer.step()
train_acc += np.sum(np.argmax(train_pred.cpu().data.numpy(), axis=1) == data[1].numpy())
train_loss += batch_loss.item()
model.eval()
val_TP = 1.0
val_TN = 1.0
val_FN = 1.0
val_FP = 1.0
predict_total = []
label_total = []
for i, data in enumerate(val_loader):
val_pred = model(data[0])
batch_loss = loss(val_pred, data[1])
predict_val = np.argmax(val_pred.cpu().data.numpy(), axis=1)
predict_total = np.append(predict_total, predict_val)
label_val = data[1].numpy()
label_total = np.append(label_total, label_val)
val_acc += np.sum(np.argmax(val_pred.cpu().data.numpy(), axis=1) == data[1].numpy())
val_loss += batch_loss.item()
val_TP = ((predict_total == 1) & (label_total == 1)).sum().item()
val_TN = ((predict_total == 0) & (label_total == 0)).sum().item()
val_FN = ((predict_total == 0) & (label_total == 1)).sum().item()
val_FP = ((predict_total == 1) & (label_total == 0)).sum().item()
val_spe = val_TN/(val_FP + val_TN + 0.001)
val_rec = val_TP/(val_TP + val_FN + 0.001)
test_acc = (val_TP+val_TN)/(val_FP + val_TN + val_TP + val_FN + 0.001)
val_acc = val_acc / val_set.__len__()
print('%3.6f %3.6f %3.6f %3.6f' % (train_acc / train_set.__len__(), train_loss, val_acc, val_loss))
if (val_acc > best_acc):
with open('save/AET_V_0.txt', 'w') as f:
f.write(str(epoch) + '\t' + str(val_acc) + '\t' + str(val_spe) + '\t' + str(val_rec) + '\n')
torch.save(model.state_dict(), 'save/model.pth')
best_acc = val_acc
for name, param in model.named_parameters():
if param.requires_grad:
print(param[0])
if __name__ == '__main__':
main()
Datasets
Models
