yuzhenhua
5 years ago
parent
1f9ff76371
commit
4128fb1155
@ -0,0 +1,68 @@
|
|||||||
|
# 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.
|
||||||
|
# ============================================================================
|
||||||
|
"""ncf export file"""
|
||||||
|
import argparse
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from mindspore import Tensor, context, load_checkpoint, load_param_into_net, export
|
||||||
|
|
||||||
|
import src.constants as rconst
|
||||||
|
from src.config import cfg
|
||||||
|
from ncf import NCFModel, PredictWithSigmoid
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser(description='ncf export')
|
||||||
|
parser.add_argument("--device_id", type=int, default=0, help="Device id")
|
||||||
|
parser.add_argument("--ckpt_file", type=str, required=True, help="Checkpoint file path.")
|
||||||
|
parser.add_argument("--dataset", type=str, default="ml-1m", choices=["ml-1m", "ml-20m"], help="Dataset.")
|
||||||
|
parser.add_argument("--file_name", type=str, default="ncf", help="output file name.")
|
||||||
|
parser.add_argument('--file_format', type=str, choices=["AIR", "ONNX", "MINDIR"], default='AIR', help='file format')
|
||||||
|
parser.add_argument("--device_target", type=str, default="Ascend",
|
||||||
|
choices=["Ascend", "GPU", "CPU"], help="device target (default: Ascend)")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
topk = rconst.TOP_K
|
||||||
|
num_eval_neg = rconst.NUM_EVAL_NEGATIVES
|
||||||
|
|
||||||
|
if args.dataset == "ml-1m":
|
||||||
|
num_eval_users = 6040
|
||||||
|
num_eval_items = 3706
|
||||||
|
elif args.dataset == "ml-20m":
|
||||||
|
num_eval_users = 138493
|
||||||
|
num_eval_items = 26744
|
||||||
|
else:
|
||||||
|
raise ValueError("not supported dataset")
|
||||||
|
|
||||||
|
ncf_net = NCFModel(num_users=num_eval_users,
|
||||||
|
num_items=num_eval_items,
|
||||||
|
num_factors=cfg.num_factors,
|
||||||
|
model_layers=cfg.layers,
|
||||||
|
mf_regularization=0,
|
||||||
|
mlp_reg_layers=[0.0, 0.0, 0.0, 0.0],
|
||||||
|
mf_dim=16)
|
||||||
|
|
||||||
|
param_dict = load_checkpoint(args.ckpt_file)
|
||||||
|
load_param_into_net(ncf_net, param_dict)
|
||||||
|
|
||||||
|
network = PredictWithSigmoid(ncf_net, topk, num_eval_neg)
|
||||||
|
|
||||||
|
users = Tensor(np.zeros([cfg.eval_batch_size, 1]).astype(np.int32))
|
||||||
|
items = Tensor(np.zeros([cfg.eval_batch_size, 1]).astype(np.int32))
|
||||||
|
masks = Tensor(np.zeros([cfg.eval_batch_size, 1]).astype(np.float32))
|
||||||
|
|
||||||
|
input_data = [users, items, masks]
|
||||||
|
export(network, *input_data, file_name=args.file_name, file_format=args.file_format)
|
||||||
@ -1,112 +0,0 @@
|
|||||||
# 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 NCF air file."""
|
|
||||||
import os
|
|
||||||
import argparse
|
|
||||||
from absl import logging
|
|
||||||
|
|
||||||
from mindspore.train.serialization import load_checkpoint, load_param_into_net, export
|
|
||||||
from mindspore import Tensor, context, Model
|
|
||||||
|
|
||||||
import constants as rconst
|
|
||||||
from dataset import create_dataset
|
|
||||||
from metrics import NCFMetric
|
|
||||||
from ncf import NCFModel, NetWithLossClass, TrainStepWrap, PredictWithSigmoid
|
|
||||||
|
|
||||||
|
|
||||||
logging.set_verbosity(logging.INFO)
|
|
||||||
|
|
||||||
|
|
||||||
def argparse_init():
|
|
||||||
"""Argparse init method"""
|
|
||||||
parser = argparse.ArgumentParser(description='NCF')
|
|
||||||
|
|
||||||
parser.add_argument("--data_path", type=str, default="./dataset/") # The location of the input data.
|
|
||||||
parser.add_argument("--dataset", type=str, default="ml-1m", choices=["ml-1m", "ml-20m"]) # Dataset to be trained and evaluated. ["ml-1m", "ml-20m"]
|
|
||||||
parser.add_argument("--eval_batch_size", type=int, default=160000) # The batch size used for evaluation.
|
|
||||||
parser.add_argument("--layers", type=int, default=[64, 32, 16]) # The sizes of hidden layers for MLP
|
|
||||||
parser.add_argument("--num_factors", type=int, default=16) # The Embedding size of MF model.
|
|
||||||
parser.add_argument("--output_path", type=str, default="./output/") # The location of the output file.
|
|
||||||
parser.add_argument("--eval_file_name", type=str, default="eval.log") # Eval output file.
|
|
||||||
parser.add_argument("--checkpoint_file_path", type=str, default="./checkpoint/NCF.ckpt") # The location of the checkpoint file.
|
|
||||||
return parser
|
|
||||||
|
|
||||||
|
|
||||||
def export_air_file():
|
|
||||||
""""Export file for eval"""
|
|
||||||
parser = argparse_init()
|
|
||||||
args, _ = parser.parse_known_args()
|
|
||||||
|
|
||||||
if not os.path.exists(args.output_path):
|
|
||||||
os.makedirs(args.output_path)
|
|
||||||
|
|
||||||
layers = args.layers
|
|
||||||
num_factors = args.num_factors
|
|
||||||
topk = rconst.TOP_K
|
|
||||||
num_eval_neg = rconst.NUM_EVAL_NEGATIVES
|
|
||||||
|
|
||||||
ds_eval, num_eval_users, num_eval_items = create_dataset(test_train=False, data_dir=args.data_path,
|
|
||||||
dataset=args.dataset, train_epochs=0,
|
|
||||||
eval_batch_size=args.eval_batch_size)
|
|
||||||
print("ds_eval.size: {}".format(ds_eval.get_dataset_size()))
|
|
||||||
|
|
||||||
ncf_net = NCFModel(num_users=num_eval_users,
|
|
||||||
num_items=num_eval_items,
|
|
||||||
num_factors=num_factors,
|
|
||||||
model_layers=layers,
|
|
||||||
mf_regularization=0,
|
|
||||||
mlp_reg_layers=[0.0, 0.0, 0.0, 0.0],
|
|
||||||
mf_dim=16)
|
|
||||||
param_dict = load_checkpoint(args.checkpoint_file_path)
|
|
||||||
load_param_into_net(ncf_net, param_dict)
|
|
||||||
|
|
||||||
loss_net = NetWithLossClass(ncf_net)
|
|
||||||
train_net = TrainStepWrap(loss_net)
|
|
||||||
train_net.set_train()
|
|
||||||
eval_net = PredictWithSigmoid(ncf_net, topk, num_eval_neg)
|
|
||||||
|
|
||||||
ncf_metric = NCFMetric()
|
|
||||||
model = Model(train_net, eval_network=eval_net, metrics={"ncf": ncf_metric})
|
|
||||||
|
|
||||||
ncf_metric.clear()
|
|
||||||
out = model.eval(ds_eval)
|
|
||||||
|
|
||||||
eval_file_path = os.path.join(args.output_path, args.eval_file_name)
|
|
||||||
eval_file = open(eval_file_path, "a+")
|
|
||||||
eval_file.write("EvalCallBack: HR = {}, NDCG = {}\n".format(out['ncf'][0], out['ncf'][1]))
|
|
||||||
eval_file.close()
|
|
||||||
print("EvalCallBack: HR = {}, NDCG = {}".format(out['ncf'][0], out['ncf'][1]))
|
|
||||||
|
|
||||||
param_dict = load_checkpoint(args.checkpoint_file_path)
|
|
||||||
# load the parameter into net
|
|
||||||
load_param_into_net(eval_net, param_dict)
|
|
||||||
|
|
||||||
input_tensor_list = []
|
|
||||||
for data in ds_eval:
|
|
||||||
for j in data:
|
|
||||||
input_tensor_list.append(Tensor(j))
|
|
||||||
print(len(a))
|
|
||||||
break
|
|
||||||
print(input_tensor_list)
|
|
||||||
export(eval_net, *input_tensor_list, file_name='NCF.air', file_format='AIR')
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
devid = int(os.getenv('DEVICE_ID'))
|
|
||||||
context.set_context(mode=context.GRAPH_MODE,
|
|
||||||
device_target="Davinci",
|
|
||||||
save_graphs=True,
|
|
||||||
device_id=devid)
|
|
||||||
|
|
||||||
export_air_file()
|
|
||||||
Loading…
Reference in new issue