Browse Source

feat: train code added

master
Mohammad Hashemi 2 years ago
parent
commit
886e614e75
1 changed files with 108 additions and 0 deletions
  1. 108
    0
      train.py

+ 108
- 0
train.py View File

@@ -0,0 +1,108 @@
import os
import argparse
import torch
import torch.nn as nn
import numpy as np
import pickle
from wd_gc import WD_GCN
from data_loader import DataLoader
from utils import f1_score

parser = argparse.ArgumentParser()
# directories
parser.add_argument('--data_root', type=str, default='data/', help='path to the root of data directory')
parser.add_argument('--dataset', type=str, default='bitcoin_alpha', help='dataset name')
parser.add_argument('--dataset_mat_file_path', type=str, default='Bitcoin_Alpha/saved_content_bitcoin_alpha.mat', help='dataset mat file path')
parser.add_argument('--save_dir', type=str, default='results_edge_classification_bitcoin_alpha/', help='path to save directory')
# model params
parser.add_argument('--hidden_feat', type=list, default=[6, 2], help='hidden feature sizes')
parser.add_argument('--time_partitioning', type=list, help='time partitioning for train/test/val split')
# training params
parser.add_argument('--epochs', type=int, default=1001, help='number of epochs')
parser.add_argument('--optimizer', type=str, default='SGD', help='optimizer')
parser.add_argument('--learning_rate', type=float, default=0.01, help='learning rate')
parser.add_argument('--momentum', type=float, default=0.9, help='momentum')
parser.add_argument('--alpha_vec', type=list,
default=[.75, .76, .77, .78, .79, .80, .81, .82, .83, .84, .85,
.86, .87, .88, .89, .90, .91, .92, .93, .94, .95], help='alpha rate for weighted cross-entropy')
args = parser.parse_args()

print('__pyTorch VERSION:', torch.__version__)
print('# of available devices ', torch.cuda.device_count())
print('Is cuda GPU available? ', torch.cuda.is_available())
print('========================')
if torch.cuda.is_available():
dev = "cuda:1"
else:
dev = "cpu"
device = torch.device(dev)
torch.cuda.set_device(device)
print('Current cuda device', torch.cuda.current_device())

# Make the save directory
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)

time_slices = args.time_partitioning
S_train, S_val, S_test = time_slices[0], time_slices[1], time_slices[2]

# Create data loader

data_loader = DataLoader(S_train=S_train, S_val=S_val, S_test=S_test,
data_dir=args.data_root,
mat_file_path=args.dataset_mat_file_path)

# load data
data, C, targets, edges = data_loader.load_data()

# train
print("Training Started...")
for alpha in args.alpha_vec:
print(">> alpha = {}".format(alpha))
class_weights = torch.tensor([alpha, 1.0 - alpha]).to(device)
save_res_fname = "results_WDGCN" + "_w" + str(round(float(class_weights[0]*100))) + "_" + args.dataset

# model definition
gcn = WD_GCN(C['C_train'], data['X_train'].to(device), edges["edges_train"].to(device), args.hidden_feat)
gcn.to(device)

if args.optimizer == "SGD":
optimizer = torch.optim.SGD(gcn.parameters(), lr=args.learning_rate, momentum=args.momentum)
criterion = nn.CrossEntropyLoss(weight=class_weights)
ep_acc_loss = np.zeros((args.epochs, 12)) # (precision_train, recall_train, f1_train, loss_train, precision_val, recall_val, f1_val, loss_val, precision_test, recall_test, f1_test, loss_test)
for ep in range(args.epochs):
# compute loss and take step
optimizer.zero_grad()
output_train = gcn(C['C_train'], data['X_train'].to(device), edges["edges_train"].to(device)) # forward passing
loss_train = criterion(output_train, targets['target_train'].to(device))
loss_train.backward() # backward propagation
optimizer.step()

with torch.no_grad():
prediction_train = torch.argmax(output_train, dim=1)
f1_train, precision_train, recall_train = f1_score(prediction_train, targets['target_train'].to(device))
if ep % 100 == 0:
# validation
output_val = gcn(C['C_val'], data['X_val'].to(device), edges["edges_val"].to(device))
prediction_val = torch.argmax(output_val, dim=1)
f1_val, precision_val, recall_val = f1_score(prediction_val, targets['target_val'].to(device))
loss_val = criterion(output_val, targets['target_val'].to(device))

# test
output_test = gcn(C['C_test'], data['X_test'].to(device), edges["edges_test"].to(device))
prediction_test = torch.argmax(output_test, dim=1)
f1_test, precision_test, recall_test = f1_score(prediction_test, targets['target_test'].to(device))
loss_test = criterion(output_test, targets['target_test'].to(device))

# log
print("Epoch %d:" % ep)
print("%d/%d: Train precision/recall/f1 %.4f/ %.4f/ %.4f. Train loss %.4f." % (ep, args.epochs, precision_train, recall_train, f1_train, loss_train))
print("%d/%d: Val precision/recall/f1 %.4f/ %.4f/ %.4f. Val loss %.4f." % (ep, args.epochs, precision_val, recall_val, f1_val, loss_val))
print("%d/%d: Test precision/recall/f1 %.4f/ %.4f/ %.4f. Test loss %.4f." % (ep, args.epochs, precision_test, recall_test, f1_test, loss_test))

ep_acc_loss[ep] = [precision_train, recall_train, f1_train, loss_train,
precision_val, recall_val, f1_val, loss_val,
precision_test, recall_test, f1_test, loss_test]


pickle.dump(ep_acc_loss, open(args.save_dir + save_res_fname, "wb"))

Loading…
Cancel
Save