• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15
16import time
17import pytest
18import numpy as np
19from mindspore import context
20from mindspore import Tensor
21from tests.models.official.gnn.gcn.src.gcn import GCN
22from tests.models.official.gnn.gcn.src.metrics import LossAccuracyWrapper, TrainNetWrapper
23from tests.models.official.gnn.gcn.src.config import ConfigGCN
24from tests.models.official.gnn.gcn.src.dataset import get_adj_features_labels, get_mask
25
26
27DATA_DIR = '/home/workspace/mindspore_dataset/cora/cora_mr/cora_mr'
28TRAIN_NODE_NUM = 140
29EVAL_NODE_NUM = 500
30TEST_NODE_NUM = 1000
31SEED = 20
32
33
34@pytest.mark.level0
35@pytest.mark.platform_arm_ascend_training
36@pytest.mark.platform_x86_ascend_training
37@pytest.mark.env_onecard
38def test_gcn():
39    print("test_gcn begin")
40    np.random.seed(SEED)
41    context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
42    config = ConfigGCN()
43    config.dropout = 0.0
44    adj, feature, label_onehot, _ = get_adj_features_labels(DATA_DIR)
45
46    nodes_num = label_onehot.shape[0]
47    train_mask = get_mask(nodes_num, 0, TRAIN_NODE_NUM)
48    eval_mask = get_mask(nodes_num, TRAIN_NODE_NUM, TRAIN_NODE_NUM + EVAL_NODE_NUM)
49    test_mask = get_mask(nodes_num, nodes_num - TEST_NODE_NUM, nodes_num)
50
51    class_num = label_onehot.shape[1]
52    input_dim = feature.shape[1]
53    gcn_net = GCN(config, input_dim, class_num)
54    gcn_net.add_flags_recursive(fp16=True)
55
56    adj = Tensor(adj)
57    feature = Tensor(feature)
58
59    eval_net = LossAccuracyWrapper(gcn_net, label_onehot, eval_mask, config.weight_decay)
60    test_net = LossAccuracyWrapper(gcn_net, label_onehot, test_mask, config.weight_decay)
61    train_net = TrainNetWrapper(gcn_net, label_onehot, train_mask, config)
62
63    loss_list = []
64    for epoch in range(config.epochs):
65        t = time.time()
66
67        train_net.set_train()
68        train_result = train_net(adj, feature)
69        train_loss = train_result[0].asnumpy()
70        train_accuracy = train_result[1].asnumpy()
71
72        eval_net.set_train(False)
73        eval_result = eval_net(adj, feature)
74        eval_loss = eval_result[0].asnumpy()
75        eval_accuracy = eval_result[1].asnumpy()
76
77        loss_list.append(eval_loss)
78        print("Epoch:", '%04d' % (epoch + 1), "train_loss=", "{:.5f}".format(train_loss),
79              "train_acc=", "{:.5f}".format(train_accuracy), "val_loss=", "{:.5f}".format(eval_loss),
80              "val_acc=", "{:.5f}".format(eval_accuracy), "time=", "{:.5f}".format(time.time() - t))
81
82        if epoch > config.early_stopping and loss_list[-1] > np.mean(loss_list[-(config.early_stopping+1):-1]):
83            print("Early stopping...")
84            break
85
86    test_net.set_train(False)
87    test_result = test_net(adj, feature)
88    test_loss = test_result[0].asnumpy()
89    test_accuracy = test_result[1].asnumpy()
90    print("Test set results:", "loss=", "{:.5f}".format(test_loss),
91          "accuracy=", "{:.5f}".format(test_accuracy))
92    assert test_accuracy > 0.812
93