summaryrefslogtreecommitdiffstats
path: root/lib/oeqa/runtime/miutils/tests/dldt_model_optimizer_test.py
blob: 7d3db15bbc13214d072c5359a5e620644b448659 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import os

class DldtModelOptimizerTest(object):
    mo_input_files = {'model': 'squeezenet_v1.1.caffemodel',
                      'prototxt': 'deploy.prototxt'}
    mo_exe = 'mo.py'

    def __init__(self, target, work_dir):
        self.target = target
        self.work_dir = work_dir

    def setup(self):
        self.target.run('mkdir -p %s' % self.work_dir)

    def tear_down(self):
        self.target.run('rm -rf %s' % self.work_dir)

    def test_dldt_mo_can_create_ir(self, mo_exe_dir, mo_files_dir):
        return self.target.run('python3 %s --input_model %s --input_proto %s --output_dir %s --data_type FP16' %
                               (os.path.join(mo_exe_dir, self.mo_exe),
                                os.path.join(mo_files_dir, self.mo_input_files['model']),
                                os.path.join(mo_files_dir, self.mo_input_files['prototxt']),
                                self.work_dir))