|
- # Copyright 2022 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.
- # ============================================================================
- """run export"""
- import argparse
- import os
-
- import numpy as np
- from mindspore import Tensor
- from mindspore import context
- from mindspore import dtype as mstype
- from mindspore import load_checkpoint
- from mindspore import load_param_into_net
- from mindspore.train.serialization import export
-
- from src.utils import get_config
- from src.utils import get_model_dataset
- from src.utils import get_params_for_net
-
-
- def run_export(args):
- """run export"""
- cfg_path = args.cfg_path
- ckpt_path = args.ckpt_path
- file_name = args.file_name
- file_format = args.file_format
-
- cfg = get_config(cfg_path)
-
- device_target = cfg['train_config']['device_target']
- device_id = int(os.getenv('DEVICE_ID', '0'))
- context.set_context(mode=context.GRAPH_MODE, device_target=device_target, device_id=device_id)
-
- pointpillarsnet, _ = get_model_dataset(cfg)
-
- params = load_checkpoint(ckpt_path)
- new_params = get_params_for_net(params)
- load_param_into_net(pointpillarsnet, new_params)
-
- v = cfg['eval_input_reader']['max_number_of_voxels']
- p = cfg['model']['voxel_generator']['max_number_of_points_per_voxel']
- n = cfg['model']['num_point_features']
- voxels = Tensor(np.zeros((1, v, p, n)), mstype.float32)
- num_points = Tensor(np.zeros((1, v)), mstype.int32)
- coors = Tensor(np.zeros((1, v, 4)), mstype.int32)
- if cfg['model']['use_bev']:
- pc_range = np.array(cfg['model']['voxel_generator']['point_cloud_range'])
- voxel_size = np.array(cfg['model']['voxel_generator']['voxel_size'])
- x, y, z = ((pc_range[3:] - pc_range[:3]) / voxel_size).astype('int32')
- bev_map = Tensor(np.zeros((1,) + (z, x * 2, y * 2)), mstype.float32)
-
- export(pointpillarsnet, voxels, num_points, coors, bev_map, file_name=file_name, file_format=file_format)
- else:
- export(pointpillarsnet, voxels, num_points, coors, file_name=file_name, file_format=file_format)
-
- print(f'{file_name}.mindir exported successfully!')
-
-
- if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument('--cfg_path', required=True, help='')
- parser.add_argument('--ckpt_path', required=True, help='')
- parser.add_argument('--file_name', default='model', help='')
- parser.add_argument('--file_format', default='MINDIR', choices=['MINDIR', 'AIR'], help='')
-
- parse_args = parser.parse_args()
- run_export(parse_args)
|