sign_trainer.py

import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# os.environ['TL_BACKEND'] = 'torch'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 
# 0:Output all; 1:Filter out INFO; 2:Filter out INFO and WARNING; 3:Filter out INFO, WARNING, and ERROR

import argparse
import tensorlayerx as tlx
from tensorlayerx.model import TrainOneStep, WithLoss
from gammagl.datasets.flickr import Flickr
from gammagl.models.sign import SignModel
import gammagl.transforms as T

class SemiSpvzLoss(WithLoss):
    def __init__(self, net, loss_fn):
        super(SemiSpvzLoss, self).__init__(backbone=net, loss_fn=loss_fn)

    def forward(self, data, label):
        logits = self._backbone(data['xs'])
        # CrossEntropyLoss equal nll_loss(log_sigmod(h), target)
        loss = self._loss_fn(logits[data['train_mask']], label[data['train_mask']])
        return loss


def evaluate(net, xs, label, mask, metrics):
    net.set_eval()
    logits = net(xs)
    metrics.update(logits[mask], label[mask])
    acc = metrics.result()
    metrics.reset()
    return acc


def main(args):
    transform = T.Compose([T.NormalizeFeatures(), T.SIGN(args.K)])
    dataset = Flickr(args.dataset_path, transform=transform)
    graph = dataset[0]
    graph = graph.tensor()  # transform may modify some data into Numpy.ndarray.

    xs = [graph.x]
    xs += [graph[f'x{i}']for i in range(1, args.K + 1)]

    net = SignModel(K=args.K, in_feat=dataset.num_node_features,
                    hid_feat=args.hidden_dim, num_classes=dataset.num_classes,
                    drop=1 - args.keep_rate)

    optimizer = tlx.optimizers.Adam(args.lr, weight_decay=args.l2_coef)
    metrics = tlx.metrics.Accuracy()
    train_weights = net.trainable_weights

    loss_func = SemiSpvzLoss(net, tlx.losses.softmax_cross_entropy_with_logits)
    train_one_step = TrainOneStep(loss_func, optimizer, train_weights)

    best_val_acc = 0
    data = {"xs": xs,
            "train_mask": graph.train_mask}
    for epoch in range(args.n_epoch):
        net.set_train()
        train_loss = train_one_step(data, graph.y)
        val_acc = evaluate(net, xs, graph.y, graph.val_mask, metrics)
        print("Epoch [{:0>3d}] ".format(epoch + 1) \
              + "  train loss: {:.4f}".format(train_loss.item()) \
              + "  val acc: {:.4f}".format(val_acc))
        # save best model on evaluation set
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            net.save_weights(args.best_model_path + "SIGN.npz")
    net.load_weights(args.best_model_path + "SIGN.npz")
    test_acc = evaluate(net, xs, graph.y, graph.test_mask, metrics)
    print("Test acc:  {:.4f}".format(test_acc))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--lr", type=float, default=0.01, help="learnin rate")
    parser.add_argument("--dataset_path", type=str, default=r'', help="path to save dataset")
    parser.add_argument('--dataset', type=str, default='Flickr', help='Only flickr')
    parser.add_argument("--hidden_dim", type=int, default=1024, help="dimention of hidden layers")
    parser.add_argument("--best_model_path", type=str, default=r'./', help="path to save best model")
    parser.add_argument("--n_epoch", type=int, default=200, help="number of epoch")
    parser.add_argument("--l2_coef", type=float, default=0., help="l2 loss coeficient")
    parser.add_argument("--K", type=int, default=2)
    parser.add_argument("--keep_rate", type=float, default=0.5, help="keep_rate = 1 - drop_rate")
    parser.add_argument("--gpu", type=int, default=0)

    args = parser.parse_args()
    if args.gpu >= 0:
        tlx.set_device("GPU", args.gpu)
    else:
        tlx.set_device("CPU")
    main(args)