Gist to help recreate issue with GCN layers (when norm set to True)

gist.py

from dgl.nn.pytorch import GraphConv
import torch.nn as nn
import torch.optim as optim
import torch as ch
from tqdm import tqdm
from botdet.data.dataset_botnet import BotnetDataset
from botdet.data.dataloader import GraphDataLoader
import argparse


LOCAL_DATA_DIR = "/localtmp/as9rw/datasets/botnet/chord"


class GCN(nn.Module):
    def __init__(self, n_hidden, n_layers, n_inp=1, n_classes=2, residual=False, norm='both', dropout=0.5):
        super(GCN, self).__init__()
        self.layers = nn.ModuleList()
        self.residual = residual
        self.drop_p = dropout

        # input layer
        self.layers.append(
            GraphConv(n_inp, n_hidden, norm=norm))

        # hidden layers
        for i in range(n_layers-1):
            self.layers.append(
                GraphConv(n_hidden, n_hidden, norm=norm))

        # output layer
        self.final = GraphConv(n_hidden, n_classes, norm=norm)
        self.activation = nn.ReLU()
        if self.drop_p > 0:
            self.dropout = nn.Dropout(p=dropout)

    def forward(self, g, latent=None):

        if latent is not None:
            if latent < 0 or latent > len(self.layers):
                raise ValueError("Invald interal layer requested")

        x = g.ndata['x']
        for i, layer in enumerate(self.layers):
            xo = self.activation(layer(g, x))

            if self.drop_p > 0:
                xo = self.dropout(xo)

            # Add prev layer directly, if requested
            if self.residual and i != 0:
                xo = self.activation(xo + x)

            x = xo

            # Return representation, if requested
            if i == latent:
                return x

        return self.final(g, x)


def true_positive(pred, target):
    return (target[pred == 1] == 1).sum().item()


def get_metrics(y, y_pred, threshold=0.5):
    y_ = 1 * (y_pred > threshold)
    tp = true_positive(y_, y)
    precision = tp / ch.sum(y_ == 1)
    recall = tp / ch.sum(y == 1)
    f1 = (2 * precision * recall) / (precision + recall)

    precision = precision.item()
    recall = recall.item()
    f1 = f1.item()

    # Check for NaNs
    if precision != precision:
        precision = 0
    if recall != recall:
        recall = 0
    if f1 != f1:
        f1 = 0

    return (precision, recall, f1)


def epoch(model, loader, gpu, optimizer=None, verbose=False):
    loss_func = nn.CrossEntropyLoss()
    is_train = True
    if optimizer is None:
        is_train = False

    tot_loss, precision, recall, f1 = 0, 0, 0, 0
    iterator = enumerate(loader)
    if verbose:
        iterator = tqdm(iterator, total=len(loader))

    for e, batch in iterator:

        # Clear gradients
        if is_train:
            optimizer.zero_grad()

        if gpu:
            # Shift graph to GPU
            batch = batch.to('cuda')

        # Get model predictions and get loss
        labels = batch.ndata['y'].long()
        logits = model(batch)
        loss = loss_func(logits, labels)

        # Logit values on botnet nodes
        # If model is learning, first number should be greater than second
        avg_o_logits = ch.mean(logits[labels == 1, 1]).detach().item()
        avg_z_logits = ch.mean(logits[labels == 1, 0]).detach().item()

        probs = ch.softmax(logits, dim=1)[:, 1]

        # Backprop gradients if training
        if is_train:
            loss.backward()
            optimizer.step()

        # Get metrics
        m = get_metrics(labels, probs)
        precision += m[0]
        recall += m[1]
        f1 += m[2]

        tot_loss += loss.detach().item()
        if verbose:
            iterator.set_description(
                "Loss: %.5f | Precision: %.3f | Recall: %.3f | F-1: %.3f | Avg-logits(0) : %.3f | Avg-logits(1) : %.3f" %
                (tot_loss / (e+1), precision / (e+1), recall / (e+1), f1 / (e+1), avg_z_logits, avg_o_logits))
    return tot_loss / (e+1)


def train_model(net, ds, args):
    train_loader, test_loader = ds.get_loaders(1, shuffle=False)
    optimizer = optim.Adam(net.parameters(), lr=args.lr)

    for e in range(args.epochs):
        # Train
        print("[Train]")
        net.train()
        epoch(net, train_loader, args.gpu, optimizer, verbose=True)

        # Test on train data (results for this part make no sense)
        print("[Eval on train data]")
        net.eval()

        epoch(net, train_loader, args.gpu, None, verbose=True)

        # Test on test data
        print("[Eval on train data]")
        net.eval()

        epoch(net, test_loader, args.gpu, None, verbose=True)

        print()


class BotNetWrapper:
    def __init__(self, dataset_name="chord"):
        self.dataset_name = dataset_name

        # For the sake of this GIST, use val data for training
        # And test data for testing
        # To make experimentation faster
        # Anyway, the issue is with model having difference in performance
        # In .train() and .eval() modes

        self.train_data = BotnetDataset(
            name="chord", root=LOCAL_DATA_DIR,
            in_memory=True, split='val',
            graph_format='dgl')

        self.test_data = BotnetDataset(
            name="chord", root=LOCAL_DATA_DIR,
            in_memory=True, split='test',
            graph_format='dgl')

    def get_loaders(self, batch_size, shuffle=False, num_workers=0):
        # Create loaders
        self.train_loader = GraphDataLoader(
            self.train_data, batch_size, shuffle, num_workers)
        self.test_loader = GraphDataLoader(
            self.test_data, batch_size, False, num_workers)

        return self.train_loader, self.test_loader


def main(args):
    # Make dataset
    ds = BotNetWrapper()

    # Define model
    model = GCN(n_hidden=args.hidden_channels,
                n_layers=args.num_layers, dropout=args.dropout,
                residual=True, norm=args.norm)

    if args.gpu:
        model.cuda()

    train_model(model, ds, args)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='CAIDA_Botnet')
    parser.add_argument('--norm', default='both', choices=['right', 'both'])
    parser.add_argument('--num_layers', type=int, default=6)
    parser.add_argument('--batch_size', type=int, default=1)
    parser.add_argument('--hidden_channels', type=int, default=32)
    parser.add_argument('--dropout', type=float, default=0.5)
    parser.add_argument('--lr', type=float, default=0.005)
    parser.add_argument('--gpu', action="store_true", help="Use CUDA?")
    parser.add_argument('--epochs', type=int, default=10)
    parser.add_argument("--savepath", help="path to save trained model")
    args = parser.parse_args()
    print(args)

    main(args)