#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2021/8/8 16:43
# @Author : Li Xiao
# @File : GCN_run.py
import numpy as np
import pandas as pd
import argparse
import glob
import os
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import f1_score
import torch
import torch.nn.functional as F
from gcn_model import GCN
from utils import load_data
from utils import accuracy
def setup_seed(seed):
torch.manual_seed(seed)
np.random.seed(seed)
def train(epoch, optimizer, features, adj, labels, idx_train):
'''
:param epoch: training epochs
:param optimizer: training optimizer, Adam optimizer
:param features: the omics features
:param adj: the laplace adjacency matrix
:param labels: sample labels
:param idx_train: the index of trained samples
'''
labels.to(device)
GCN_model.train()
optimizer.zero_grad()
output = GCN_model(features, adj)
loss_train = F.cross_entropy(output[idx_train], labels[idx_train])
acc_train = accuracy(output[idx_train], labels[idx_train])
loss_train.backward()
optimizer.step()
if (epoch+1) % 10 ==0:
print('Epoch: %.2f | loss train: %.4f | acc train: %.4f' %(epoch+1, loss_train.item(), acc_train.item()))
return loss_train.data.item()
def test(features, adj, labels, idx_test):
'''
:param features: the omics features
:param adj: the laplace adjacency matrix
:param labels: sample labels
:param idx_test: the index of tested samples
'''
GCN_model.eval()
output = GCN_model(features, adj)
loss_test = F.cross_entropy(output[idx_test], labels[idx_test])
#calculate the accuracy
acc_test = accuracy(output[idx_test], labels[idx_test])
#output is the one-hot label
ot = output[idx_test].detach().cpu().numpy()
#change one-hot label to digit label
ot = np.argmax(ot, axis=1)
#original label
lb = labels[idx_test].detach().cpu().numpy()
print('predict label: ', ot)
print('original label: ', lb)
#calculate the f1 score
f = f1_score(ot, lb, average='weighted')
print("Test set results:",
"loss= {:.4f}".format(loss_test.item()),
"accuracy= {:.4f}".format(acc_test.item()))
#return accuracy and f1 score
return acc_test.item(), f
def predict(features, adj, sample, idx):
'''
:param features: the omics features
:param adj: the laplace adjacency matrix
:param sample: all sample names
:param idx: the index of predict samples
:return:
'''
GCN_model.eval()
output = GCN_model(features, adj)
predict_label = output.detach().cpu().numpy()
predict_label = np.argmax(predict_label, axis=1).tolist()
#print(predict_label)
res_data = pd.DataFrame({'Sample':sample, 'predict_label':predict_label})
res_data = res_data.iloc[idx,:]
#print(res_data)
res_data.to_csv('result/GCN_predicted_data.csv', header=True, index=False)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--featuredata', '-fd', type=str, required=True, help='The vector feature file.')
parser.add_argument('--adjdata', '-ad', type=str, required=True, help='The adjacency matrix file.')
parser.add_argument('--labeldata', '-ld', type=str, required=True, help='The sample label file.')
parser.add_argument('--testsample', '-ts', type=str, help='Test sample names file.')
parser.add_argument('--mode', '-m', type=int, choices=[0,1], default=0,
help='mode 0: 10-fold cross validation; mode 1: train and test a model.')
parser.add_argument('--seed', '-s', type=int, default=0, help='Random seed, default=0.')
parser.add_argument('--device', '-d', type=str, choices=['cpu', 'gpu'], default='cpu',
help='Training on cpu or gpu, default: cpu.')
parser.add_argument('--epochs', '-e', type=int, default=150, help='Training epochs, default: 150.')
parser.add_argument('--learningrate', '-lr', type=float, default=0.001, help='Learning rate, default: 0.001.')
parser.add_argument('--weight_decay', '-w', type=float, default=0.01,
help='Weight decay (L2 loss on parameters), methods to avoid overfitting, default: 0.01')
parser.add_argument('--hidden', '-hd',type=int, default=64, help='Hidden layer dimension, default: 64.')
parser.add_argument('--dropout', '-dp', type=float, default=0.5, help='Dropout rate, methods to avoid overfitting, default: 0.5.')
parser.add_argument('--threshold', '-t', type=float, default=0.005, help='Threshold to filter edges, default: 0.005')
parser.add_argument('--nclass', '-nc', type=int, default=4, help='Number of classes, default: 4')
parser.add_argument('--patience', '-p', type=int, default=20, help='Patience')
args = parser.parse_args()
# Check whether GPUs are available
device = torch.device('cpu')
if args.device == 'gpu':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# set random seed
setup_seed(args.seed)
# load input files
adj, data, label = load_data(args.adjdata, args.featuredata, args.labeldata, args.threshold)
# change dataframe to Tensor
adj = torch.tensor(adj, dtype=torch.float, device=device)
features = torch.tensor(data.iloc[:, 1:].values, dtype=torch.float, device=device)
labels = torch.tensor(label.iloc[:, 1].values, dtype=torch.long, device=device)
print('Begin training model...')
# 10-fold cross validation
if args.mode == 0:
skf = StratifiedKFold(n_splits=10, shuffle=True)
acc_res, f1_res = [], [] #record accuracy and f1 score
# split train and test data
for idx_train, idx_test in skf.split(data.iloc[:, 1:], label.iloc[:, 1]):
# initialize a model
GCN_model = GCN(n_in=features.shape[1], n_hid=args.hidden, n_out=args.nclass, dropout=args.dropout)
GCN_model.to(device)
# define the optimizer
optimizer = torch.optim.Adam(GCN_model.parameters(), lr=args.learningrate, weight_decay=args.weight_decay)
idx_train, idx_test= torch.tensor(idx_train, dtype=torch.long, device=device), torch.tensor(idx_test, dtype=torch.long, device=device)
for epoch in range(args.epochs):
train(epoch, optimizer, features, adj, labels, idx_train)
# calculate the accuracy and f1 score
ac, f1= test(features, adj, labels, idx_test)
acc_res.append(ac)
f1_res.append(f1)
print('10-fold Acc(%.4f, %.4f) F1(%.4f, %.4f)' % (np.mean(acc_res), np.std(acc_res), np.mean(f1_res), np.std(f1_res)))
#predict(features, adj, data['Sample'].tolist(), data.index.tolist())
elif args.mode == 1:
# load test samples
test_sample_df = pd.read_csv(args.testsample, header=0, index_col=None)
test_sample = test_sample_df.iloc[:, 0].tolist()
all_sample = data['Sample'].tolist()
train_sample = list(set(all_sample)-set(test_sample))
#get index of train samples and test samples
train_idx = data[data['Sample'].isin(train_sample)].index.tolist()
test_idx = data[data['Sample'].isin(test_sample)].index.tolist()
GCN_model = GCN(n_in=features.shape[1], n_hid=args.hidden, n_out=args.nclass, dropout=args.dropout)
GCN_model.to(device)
optimizer = torch.optim.Adam(GCN_model.parameters(), lr=args.learningrate, weight_decay=args.weight_decay)
idx_train, idx_test = torch.tensor(train_idx, dtype=torch.long, device=device), torch.tensor(test_idx, dtype=torch.long, device=device)
'''
save a best model (with the minimum loss value)
if the loss didn't decrease in N epochs，stop the train process.
N can be set by args.patience
'''
loss_values = [] #record the loss value of each epoch
# record the times with no loss decrease, record the best epoch
bad_counter, best_epoch = 0, 0
best = 1000 #record the lowest loss value
for epoch in range(args.epochs):
loss_values.append(train(epoch, optimizer, features, adj, labels, idx_train))
if loss_values[-1] < best:
best = loss_values[-1]
best_epoch = epoch
bad_counter = 0
else:
bad_counter += 1 #In this epoch, the loss value didn't decrease
if bad_counter == args.patience:
break
#save model of this epoch
torch.save(GCN_model.state_dict(), 'model/GCN/{}.pkl'.format(epoch))
#reserve the best model, delete other models
files = glob.glob('model/GCN/*.pkl')
for file in files:
name = file.split('\\')[1]
epoch_nb = int(name.split('.')[0])
#print(file, name, epoch_nb)
if epoch_nb != best_epoch:
os.remove(file)
print('Training finished.')
print('The best epoch model is ',best_epoch)
GCN_model.load_state_dict(torch.load('model/GCN/{}.pkl'.format(best_epoch)))
predict(features, adj, all_sample, test_idx)
print('Finished!')
Datasets
Models
