fix textcnn export on master branch

This commit is contained in:
Yanjun Peng 2020-12-28 19:44:42 +08:00
parent b35435cc02
commit a152c510c5
2 changed files with 12 additions and 5 deletions

View File

@ -51,7 +51,7 @@ def parse_args():
# distributed related
parser.add_argument('--is_distributed', type=int, default=1, help='if multi device')
parser.add_argument('--ckpt_path', type=str, default='', help='ckpt path to save')
parser.add_argument('--eval_each_epoch', type=int, default=0, help='ckpt path to save')
parser.add_argument('--eval_each_epoch', type=int, default=0, help='evaluate on each epoch')
args, _ = parser.parse_known_args()
args.image_size = config.image_size
args.num_classes = config.num_classes

View File

@ -21,9 +21,9 @@ import numpy as np
from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context
from src.config import cfg
from src.config import cfg_mr, cfg_subj, cfg_sst2
from src.textcnn import TextCNN
from src.dataset import MovieReview
from src.dataset import MovieReview, SST2, Subjectivity
parser = argparse.ArgumentParser(description='TextCNN export')
parser.add_argument("--device_id", type=int, default=0, help="device id")
@ -32,7 +32,7 @@ parser.add_argument("--file_name", type=str, default="textcnn", help="output fil
parser.add_argument('--file_format', type=str, choices=["AIR", "ONNX", "MINDIR"], default='AIR', help='file format')
parser.add_argument("--device_target", type=str, choices=["Ascend", "GPU", "CPU"], default="Ascend",
help="device target")
parser.add_argument('--dataset_name', type=str, default='MR', choices=['MR'],
parser.add_argument('--dataset', type=str, default='MR', choices=['MR', 'SUBJ', 'SST2'],
help='dataset name.')
args = parser.parse_args()
@ -41,8 +41,15 @@ context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, d
if __name__ == '__main__':
if args.dataset_name == 'MR':
if args.dataset == 'MR':
cfg = cfg_mr
instance = MovieReview(root_dir=cfg.data_path, maxlen=cfg.word_len, split=0.9)
elif args.dataset == 'SUBJ':
cfg = cfg_subj
instance = Subjectivity(root_dir=cfg.data_path, maxlen=cfg.word_len, split=0.9)
elif args.dataset == 'SST2':
cfg = cfg_sst2
instance = SST2(root_dir=cfg.data_path, maxlen=cfg.word_len, split=0.9)
else:
raise ValueError("dataset is not support.")