from src.models import define_net, load_ckpt from src.utils import switch_precision, set_context if __name__ == '__main__': args_opt = eval_parse_args() config = set_config(args_opt) backbone_net, head_net, net = define_net(config, args_opt.is_training) #load the trained checkpoint file to the net for evaluation if args_opt.head_ckpt: load_ckpt(backbone_net, args_opt.pretrain_ckpt) load_ckpt(head_net, args_opt.head_ckpt) else: load_ckpt(net, args_opt.pretrain_ckpt) set_context(config) switch_precision(net, mstype.float16, config) dataset = create_dataset(dataset_path=args_opt.dataset_path, do_train=False, config=config) step_size = dataset.get_dataset_size() if step_size == 0: raise ValueError( "The step_size of dataset is zero. Check if the images count of train dataset is more \ than batch_size in config.py") net.set_train(False) loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean') model = Model(net, loss_fn=loss, metrics={'acc'})
type=str, choices=["AIR", "ONNX", "MINDIR"], default="AIR", help="file format") parser.add_argument('--platform', type=str, default="Ascend", choices=("Ascend", "GPU", "CPU"), help='run platform, only support GPU, CPU and Ascend') args = parser.parse_args() args.is_training = False args.run_distribute = False context.set_context(mode=context.GRAPH_MODE, device_target=args.platform) if args.platform == "Ascend": context.set_context(device_id=args.device_id) if __name__ == '__main__': cfg = set_config(args) set_context(cfg) _, _, net = define_net(cfg, args.is_training) load_ckpt(net, args.ckpt_file) input_shp = [args.batch_size, 3, cfg.image_height, cfg.image_width] input_array = Tensor( np.random.uniform(-1.0, 1.0, size=input_shp).astype(np.float32)) export(net, input_array, file_name=args.file_name, file_format=args.file_format)