|
- 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 numpy
- import argparse
- import tensorlayerx as tlx
- from tensorlayerx.model import TrainOneStep, WithLoss
- from gammagl.loader import DataLoader
- from gammagl.datasets import TUDataset
- from gammagl.models import GINModel
-
- class SemiSpvzLoss(WithLoss):
- def __init__(self, net, loss_fn):
- super(SemiSpvzLoss, self).__init__(backbone=net, loss_fn=loss_fn)
-
- def forward(self, data, y):
- train_logits = self.backbone_network(data.x, data.edge_index, data.batch)
- # train_logits = self.backbone_network(data.x, data.edge_index, None, data.batch.shape[0], data.batch)
- loss = self._loss_fn(train_logits, data.y)
- return loss
-
-
-
- def main(args):
- print("loading dataset...")
-
- path = args.dataset_path
- dataset = TUDataset(path, name=args.dataset)
-
-
- dataset_unit = len(dataset) // 10
- train_dataset = dataset[2 * dataset_unit:]
- val_dataset = dataset[:dataset_unit]
- test_dataset = dataset[dataset_unit: 2 * dataset_unit]
-
- train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
- val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)
- test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)
-
- net = GINModel(in_channels=max(dataset.num_features, 1),
-
- hidden_channels=args.hidden_dim,
- out_channels=dataset.num_classes,
- num_layers=args.num_layers,
- name="GIN")
-
-
- # net = GCNModel(feature_dim=max(dataset.num_features, 1),
- # hidden_dim=args.hidden_dim,
- # num_class=dataset.num_classes,
- # drop_rate=args.drop_rate,
- # num_layers=args.num_layers,
- # name="GCN")
-
- optimizer = tlx.optimizers.Adam(lr=args.lr, weight_decay=args.l2_coef)
-
-
- 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
- for epoch in range(args.n_epoch):
- net.set_train()
-
- for data in train_loader:
- train_loss = train_one_step(data, data.y)
-
- net.set_eval()
-
- total_correct = 0
- for data in val_loader:
- val_logits = net(data.x, data.edge_index, data.batch)
- # val_logits = net(data.x, data.edge_index, None, data.batch.shape[0], data.batch)
- pred = tlx.argmax(val_logits, axis=-1)
- total_correct += int((numpy.sum(tlx.convert_to_numpy(pred == data.y).astype(int))))
-
- val_acc = total_correct / len(val_dataset)
-
- print("Epoch [{:0>3d}] ".format(epoch + 1) \
- + " train loss: {:.4f}".format(train_loss.item()) \
- + " val acc: {:.4f}".format(val_acc))
-
- if val_acc > best_val_acc:
- best_val_acc = val_acc
- net.save_weights(args.best_model_path + net.name + ".npz", format='npz_dict')
-
- net.load_weights(args.best_model_path + net.name + ".npz", format='npz_dict')
- net.set_eval()
- total_correct = 0
- for data in test_loader:
- test_logits = net(data.x, data.edge_index, data.batch)
- # test_logits = net(data.x, data.edge_index, None, data.batch.shape[0], data.batch)
- pred = tlx.argmax(test_logits, axis=-1)
- total_correct += int((numpy.sum(tlx.convert_to_numpy(pred == data['y']).astype(int))))
- test_acc = total_correct / len(test_dataset)
-
- print("Test acc: {:.4f}".format(test_acc))
-
-
-
- if __name__ == '__main__':
- # parameters setting
- parser = argparse.ArgumentParser()
-
- parser.add_argument("--lr", type=float, default=0.001, help="learning rate")
- parser.add_argument("--n_epoch", type=int, default=100, help="number of epoch")
- parser.add_argument("--hidden_dim", type=int, default=32, help="dimention of hidden layers")
- parser.add_argument("--drop_rate", type=float, default=0.1, help="drop_rate")
- parser.add_argument("--l2_coef", type=float, default=5e-4, help="l2 loss coeficient")
- parser.add_argument('--dataset', type=str, default='MUTAG', help='dataset(MUTAG/IMDB-BINARY/REDDIT-BINARY)')
- parser.add_argument("--dataset_path", type=str, default=r'', help="path to save dataset")
- parser.add_argument("--best_model_path", type=str, default=r'./', help="path to save best model")
- parser.add_argument("--self_loops", type=int, default=1, help="number of graph self-loop")
- parser.add_argument("--num_layers", type=int, default=5, help="num of gin layers")
- parser.add_argument("--batch_size", type=int, default=100, help="batch_size of the data_loader")
- 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)
|