From 9734b93b00ba2ddd8d35abd87d51f49bf0737886 Mon Sep 17 00:00:00 2001 From: yuzhenhua Date: Sat, 17 Oct 2020 10:25:18 +0800 Subject: [PATCH] export script for gcn and gat --- mindspore/ops/operations/_quant_ops.py | 16 +++---- model_zoo/official/gnn/gat/export.py | 64 ++++++++++++++++++++++++++ model_zoo/official/gnn/gcn/export.py | 53 +++++++++++++++++++++ 3 files changed, 125 insertions(+), 8 deletions(-) create mode 100644 model_zoo/official/gnn/gat/export.py create mode 100644 model_zoo/official/gnn/gcn/export.py diff --git a/mindspore/ops/operations/_quant_ops.py b/mindspore/ops/operations/_quant_ops.py index 0c8d52a802..e186ba62f8 100644 --- a/mindspore/ops/operations/_quant_ops.py +++ b/mindspore/ops/operations/_quant_ops.py @@ -210,9 +210,9 @@ class FakeQuantWithMinMaxVars(PrimitiveWithInfer): raise ValueError(f"For '{self.name}', the shape of \'min\' cannot broadcast to the shape of \'x\'.") def infer_shape(self, x_shape, min_shape, max_shape): - validator.check_integer("x rank", len(x_shape), 1, Rel.GE, self.name) + validator.check_int(len(x_shape), 1, Rel.GE, "x rank", self.name) validator.check("min shape", min_shape, "max shape", max_shape, Rel.EQ, self.name) - validator.check_integer("min shape", len(min_shape), 1, Rel.EQ, self.name) + validator.check_int(len(min_shape), 1, Rel.EQ, "min shape", self.name) self.check_broadcast(min_shape, x_shape) return x_shape @@ -273,10 +273,10 @@ class FakeQuantWithMinMaxVarsGradient(PrimitiveWithInfer): raise ValueError(f"For '{self.name}', the shape of \'min\' cannot broadcast to the shape of \'x\'.") def infer_shape(self, dout_shape, x_shape, min_shape, max_shape): - validator.check_integer("x rank", len(x_shape), 1, Rel.GE, self.name) + validator.check_int(len(x_shape), 1, Rel.GE, "x rank", self.name) validator.check("dout shape", dout_shape, "x shape", x_shape, Rel.EQ, self.name) validator.check("min shape", min_shape, "max shape", max_shape, Rel.EQ, self.name) - validator.check_integer("min shape", len(min_shape), 1, Rel.EQ, self.name) + validator.check_int(len(min_shape), 1, Rel.EQ, "min shape", self.name) self.check_broadcast(min_shape, x_shape) return x_shape, min_shape, max_shape @@ -325,9 +325,9 @@ class FakeQuantWithMinMaxVarsPerChannel(PrimitiveWithInfer): 'narrow_range', narrow_range, (bool,), self.name) def infer_shape(self, x_shape, min_shape, max_shape): - validator.check_integer("x rank", len(x_shape), 1, Rel.GE, self.name) + validator.check_int(len(x_shape), 1, Rel.GE, "x rank", self.name) validator.check("min shape", min_shape, "max shape", max_shape, Rel.EQ, self.name) - validator.check_integer("min shape", len(min_shape), 1, Rel.EQ, self.name) + validator.check_int(len(min_shape), 1, Rel.EQ, "min shape", self.name) validator.check("min shape", min_shape[0], "x shape", x_shape[-1], Rel.EQ, self.name) return x_shape @@ -382,10 +382,10 @@ class FakeQuantWithMinMaxVarsPerChannelGradient(PrimitiveWithInfer): 'narrow_range', narrow_range, (bool,), self.name) def infer_shape(self, dout_shape, x_shape, min_shape, max_shape): - validator.check_integer("x rank", len(x_shape), 1, Rel.GE, self.name) + validator.check_int(len(x_shape), 1, Rel.GE, "x rank", self.name) validator.check("dout shape", dout_shape, "x shape", x_shape, Rel.EQ, self.name) validator.check("min shape", min_shape, "max shape", max_shape, Rel.EQ, self.name) - validator.check_integer("min shape", len(min_shape), 1, Rel.EQ, self.name) + validator.check_int(len(min_shape), 1, Rel.EQ, "min shape", self.name) validator.check("min shape", min_shape[0], "x shape", x_shape[-1], Rel.EQ, self.name) return x_shape, min_shape, max_shape diff --git a/model_zoo/official/gnn/gat/export.py b/model_zoo/official/gnn/gat/export.py new file mode 100644 index 0000000000..686580773b --- /dev/null +++ b/model_zoo/official/gnn/gat/export.py @@ -0,0 +1,64 @@ +# Copyright 2020 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""export checkpoint file into air models""" +import argparse +import numpy as np + +from mindspore import Tensor, context +from mindspore.train.serialization import load_checkpoint, export + +from src.gat import GAT +from src.config import GatConfig + +context.set_context(mode=context.GRAPH_MODE, device_target="Ascend") + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='GAT_export') + parser.add_argument('--ckpt_file', type=str, default='./ckpts/gat.ckpt', help='GAT ckpt file.') + parser.add_argument('--output_file', type=str, default='gat.air', help='GAT output air name.') + parser.add_argument('--dataset', type=str, default='cora', help='GAT dataset name.') + args_opt = parser.parse_args() + + if args_opt.dataset == "citeseer": + feature_size = [1, 3312, 3703] + biases_size = [1, 3312, 3312] + num_classes = 6 + else: + feature_size = [1, 2708, 1433] + biases_size = [1, 2708, 2708] + num_classes = 7 + + hid_units = GatConfig.hid_units + n_heads = GatConfig.n_heads + + feature = np.random.uniform(0.0, 1.0, size=feature_size).astype(np.float32) + biases = np.random.uniform(0.0, 1.0, size=biases_size).astype(np.float64) + + feature_size = feature.shape[2] + num_nodes = feature.shape[1] + + gat_net = GAT(feature_size, + num_classes, + num_nodes, + hid_units, + n_heads, + attn_drop=0.0, + ftr_drop=0.0) + + gat_net.set_train(False) + load_checkpoint(args_opt.ckpt_file, net=gat_net) + gat_net.add_flags_recursive(fp16=True) + + export(gat_net, Tensor(feature), Tensor(biases), file_name=args_opt.output_file, file_format="AIR") diff --git a/model_zoo/official/gnn/gcn/export.py b/model_zoo/official/gnn/gcn/export.py new file mode 100644 index 0000000000..e730f08c5b --- /dev/null +++ b/model_zoo/official/gnn/gcn/export.py @@ -0,0 +1,53 @@ +# Copyright 2020 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""export checkpoint file into air models""" +import argparse +import numpy as np + +from mindspore import Tensor, context +from mindspore.train.serialization import load_checkpoint, export + +from src.gcn import GCN +from src.config import ConfigGCN + +context.set_context(mode=context.GRAPH_MODE, device_target="Ascend") + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='GCN_export') + parser.add_argument('--ckpt_file', type=str, default='', help='GCN ckpt file.') + parser.add_argument('--output_file', type=str, default='gcn.air', help='GCN output air name.') + parser.add_argument('--dataset', type=str, default='cora', help='GCN dataset name.') + args_opt = parser.parse_args() + + config = ConfigGCN() + + if args_opt.dataset == "cora": + input_dim = 1433 + class_num = 7 + adj = Tensor(np.zeros((2708, 2708), np.float64)) + feature = Tensor(np.zeros((2708, 1433), np.float32)) + else: + input_dim = 3703 + class_num = 6 + adj = Tensor(np.zeros((3312, 3312), np.float64)) + feature = Tensor(np.zeros((3312, 3703), np.float32)) + + gcn_net = GCN(config, input_dim, class_num) + + gcn_net.set_train(False) + load_checkpoint(args_opt.ckpt_file, net=gcn_net) + gcn_net.add_flags_recursive(fp16=True) + + export(gcn_net, adj, feature, file_name=args_opt.output_file, file_format="AIR")