feat: train code added

train.py

@@ -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"))