Xinda Huang
huangxinda
4 years ago
1 changed files with 0 additions and 54 deletions
Unified View
Diff Options
+ 0
- 54
imperative/python/test/unit/module/test_external.py
View File
| @@ -1,54 +0,0 @@ | |||||
| # -*- coding: utf-8 -*- | |||||
| # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") | |||||
| # | |||||
| # Copyright (c) 2014-2021 Megvii Inc. All rights reserved. | |||||
| # | |||||
| # Unless required by applicable law or agreed to in writing, | |||||
| # software distributed under the License is distributed on an | |||||
| # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |||||
| import io | |||||
| import os | |||||
| import platform | |||||
| import numpy as np | |||||
| import pytest | |||||
| import megengine as mge | |||||
| import megengine.utils.comp_graph_tools as cgtools | |||||
| from megengine import Tensor | |||||
| from megengine.distributed.helper import get_device_count_by_fork | |||||
| from megengine.jit import trace | |||||
| from megengine.module import Module | |||||
| from megengine.module.external import TensorrtRuntimeSubgraph | |||||
| class MyModule(Module): | |||||
| def __init__(self, data): | |||||
| from megengine.module.external import CambriconSubgraph | |||||
| super().__init__() | |||||
| self.cambricon = CambriconSubgraph(data, "subnet0", True) | |||||
| def forward(self, inputs): | |||||
| out = self.cambricon(inputs) | |||||
| return out | |||||
| @pytest.mark.skip(reason="cambricon unimplemented") | |||||
| def test_cambricon_module(): | |||||
| model = "CambriconRuntimeOprTest.MutableBatchSize.mlu" | |||||
| model = os.path.join(os.path.dirname(__file__), model) | |||||
| with open(model, "rb") as f: | |||||
| data = f.read() | |||||
| m = MyModule(data) | |||||
| inp = Tensor( | |||||
| np.random.normal((1, 64, 32, 32)).astype(np.float16), device="cambricon0" | |||||
| ) | |||||
| def inference(inps): | |||||
| pred = m(inps) | |||||
| return pred | |||||
| pred = inference([inp]) | |||||