Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit c3e3c5c

Browse files
nimin98facebook-github-bot
nimin98
authored and
facebook-github-bot
committed
Skip tests if C2/ONNX models cannot be read (#18494)
Summary: Pull Request resolved: #18494 Today we have some C2 end2end test run requiring reading model data from external filesystem (for example, Gluster and AWS). This could be a source for flaky test when the external filesystems are not reachable during the tests. In this diff, we add try/catch logic around where we download models and open model files from external system. In case such attempts fails, we will catch the excption and let the unittest skip the current test instead of failure. I also refactor the code a little bit by removing some duplicated logic on downloading and build the c2 model data. It has been duplicated in two classes and a few functions... Reviewed By: yinghai Differential Revision: D14442241 fbshipit-source-id: da8bf56c8d096efa34ca2070de5cd10a18aad70c
1 parent 30da6c7 commit c3e3c5c

5 files changed

Lines changed: 46 additions & 147 deletions

File tree

‎caffe2/python/models/download.py‎

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -147,37 +147,39 @@ def validModelName(name):
147147
return True
148148

149149
class ModelDownloader:
150+
151+
def __init__(self, model_env_name='CAFFE2_MODELS'):
152+
self.model_env_name = model_env_name
153+
150154
def _model_dir(self, model):
151155
caffe2_home = os.path.expanduser(os.getenv('CAFFE2_HOME', '~/.caffe2'))
152-
models_dir = os.getenv('CAFFE2_MODELS', os.path.join(caffe2_home, 'models'))
156+
models_dir = os.getenv(self.model_env_name, os.path.join(caffe2_home, 'models'))
153157
return os.path.join(models_dir, model)
154158

155159
def _download(self, model):
156160
model_dir = self._model_dir(model)
157161
assert not os.path.exists(model_dir)
158162
os.makedirs(model_dir)
163+
159164
for f in ['predict_net.pb', 'init_net.pb', 'value_info.json']:
160165
url = getURLFromName(model, f)
161166
dest = os.path.join(model_dir, f)
162167
try:
163-
try:
164-
downloadFromURLToFile(url, dest,
165-
show_progress=False)
166-
except TypeError:
167-
# show_progress not supported prior to
168-
# Caffe2 78c014e752a374d905ecfb465d44fa16e02a28f1
169-
# (Sep 17, 2017)
170-
downloadFromURLToFile(url, dest)
171-
except Exception as e:
172-
print("Abort: {reason}".format(reason=e))
173-
print("Cleaning up...")
168+
downloadFromURLToFile(url, dest, show_progress=False)
169+
except TypeError:
170+
# show_progress not supported prior to
171+
# Caffe2 78c014e752a374d905ecfb465d44fa16e02a28f1
172+
# (Sep 17, 2017)
173+
downloadFromURLToFile(url, dest)
174+
except Exception:
174175
deleteDirectory(model_dir)
175-
exit(1)
176+
raise
176177

177178
def get_c2_model(self, model_name):
178179
model_dir = self._model_dir(model_name)
179180
if not os.path.exists(model_dir):
180181
self._download(model_name)
182+
181183
c2_predict_pb = os.path.join(model_dir, 'predict_net.pb')
182184
c2_predict_net = caffe2_pb2.NetDef()
183185
with open(c2_predict_pb, 'rb') as f:

‎caffe2/python/onnx/test_onnxifi.py‎

Lines changed: 4 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
from onnx.helper import make_node, make_graph, make_tensor, make_tensor_value_info, make_model
1616
from caffe2.proto import caffe2_pb2
1717
from caffe2.python import core, workspace
18-
from caffe2.python.models.download import downloadFromURLToFile, getURLFromName, deleteDirectory
18+
from caffe2.python.models.download import ModelDownloader
1919
from caffe2.python.onnx.onnxifi import onnxifi_caffe2_net
2020
from caffe2.python.onnx.tests.test_utils import TestCase
2121

@@ -109,54 +109,8 @@ def test_conv_graph(self):
109109

110110

111111
class OnnxifiTransformTest(TestCase):
112-
def _model_dir(self, model):
113-
caffe2_home = os.path.expanduser(os.getenv('CAFFE2_HOME', '~/.caffe2'))
114-
models_dir = os.getenv('CAFFE2_MODELS', os.path.join(caffe2_home, 'models'))
115-
return os.path.join(models_dir, model)
116-
117-
def _download(self, model):
118-
model_dir = self._model_dir(model)
119-
assert not os.path.exists(model_dir)
120-
os.makedirs(model_dir)
121-
for f in ['predict_net.pb', 'init_net.pb', 'value_info.json']:
122-
url = getURLFromName(model, f)
123-
dest = os.path.join(model_dir, f)
124-
try:
125-
try:
126-
downloadFromURLToFile(url, dest,
127-
show_progress=False)
128-
except TypeError:
129-
# show_progress not supported prior to
130-
# Caffe2 78c014e752a374d905ecfb465d44fa16e02a28f1
131-
# (Sep 17, 2017)
132-
downloadFromURLToFile(url, dest)
133-
except Exception as e:
134-
print("Abort: {reason}".format(reason=e))
135-
print("Cleaning up...")
136-
deleteDirectory(model_dir)
137-
exit(1)
138-
139-
# TODO: we need to modulize this function
140-
def _get_c2_model(self, model_name):
141-
model_dir = self._model_dir(model_name)
142-
if not os.path.exists(model_dir):
143-
self._download(model_name)
144-
c2_predict_pb = os.path.join(model_dir, 'predict_net.pb')
145-
c2_predict_net = caffe2_pb2.NetDef()
146-
with open(c2_predict_pb, 'rb') as f:
147-
c2_predict_net.ParseFromString(f.read())
148-
c2_predict_net.name = model_name
149-
150-
c2_init_pb = os.path.join(model_dir, 'init_net.pb')
151-
c2_init_net = caffe2_pb2.NetDef()
152-
with open(c2_init_pb, 'rb') as f:
153-
c2_init_net.ParseFromString(f.read())
154-
c2_init_net.name = model_name + '_init'
155-
156-
with open(os.path.join(model_dir, 'value_info.json')) as f:
157-
value_info = json.load(f)
158-
159-
return c2_init_net, c2_predict_net, value_info
112+
def setUp(self):
113+
self.model_downloader = ModelDownloader()
160114

161115
def _add_head_tail(self, pred_net, new_head, new_tail):
162116
orig_head = pred_net.external_input[0]
@@ -187,7 +141,7 @@ def test_resnet50_core(self):
187141
N = 1
188142
repeat = 1
189143
print("Batch size: {}, repeat inference {} times".format(N, repeat))
190-
init_net, pred_net, _ = self._get_c2_model('resnet50')
144+
init_net, pred_net, _ = self.model_downloader.get_c2_model('resnet50')
191145
self._add_head_tail(pred_net, 'real_data', 'real_softmax')
192146
input_blob_dims = (N, 3, 224, 224)
193147
input_name = "real_data"

‎caffe2/python/onnx/tests/c2_ref_test.py‎

Lines changed: 17 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -23,14 +23,14 @@
2323
import caffe2.python.onnx.backend as c2
2424

2525
import numpy as np
26-
from caffe2.python.models.download import downloadFromURLToFile, getURLFromName, deleteDirectory
26+
from caffe2.python.models.download import ModelDownloader
2727

28-
from caffe2.python.onnx.tests.test_utils import DownloadingTestCase
28+
from caffe2.python.onnx.tests.test_utils import TestCase
2929

3030
import caffe2.python._import_c_extension as C
3131

3232

33-
class TestCaffe2Basic(DownloadingTestCase):
33+
class TestCaffe2Basic(TestCase):
3434
def test_dummy_name(self):
3535
g = C.DummyName()
3636
n1 = g.new_dummy_name()
@@ -726,43 +726,33 @@ def test_cast(self):
726726
self.assertSameOutputs(c2_outputs, onnx_outputs)
727727

728728

729-
class TestCaffe2End2End(DownloadingTestCase):
730-
def _model_dir(self, model):
731-
caffe2_home = os.path.expanduser(os.getenv('CAFFE2_HOME', '~/.caffe2'))
732-
models_dir = os.getenv('ONNX_MODELS', os.path.join(caffe2_home, 'models'))
733-
return os.path.join(models_dir, model)
729+
class TestCaffe2End2End(TestCase):
730+
def setUp(self):
731+
self.model_downloader = ModelDownloader('ONNX_MODELS')
734732

735733
def _test_net(self,
736734
net_name,
737735
input_blob_dims=(1, 3, 224, 224),
738736
decimal=7):
739737
np.random.seed(seed=0)
740-
model_dir = self._model_dir(net_name)
741-
if not os.path.exists(model_dir):
742-
self._download(net_name)
743-
c2_predict_pb = os.path.join(model_dir, 'predict_net.pb')
744-
c2_predict_net = caffe2_pb2.NetDef()
745-
with open(c2_predict_pb, 'rb') as f:
746-
c2_predict_net.ParseFromString(f.read())
747-
c2_predict_net.name = net_name
748-
749-
c2_init_pb = os.path.join(model_dir, 'init_net.pb')
750-
c2_init_net = caffe2_pb2.NetDef()
751-
with open(c2_init_pb, 'rb') as f:
752-
c2_init_net.ParseFromString(f.read())
753-
c2_init_net.name = net_name + '_init'
738+
try:
739+
c2_init_net, c2_predict_net, value_info = self.model_downloader.get_c2_model(net_name)
740+
except (OSError, IOError) as e:
741+
# catch IOError/OSError that is caused by FileNotFoundError and PermissionError
742+
self.skipTest(str(e))
754743

744+
# start to run the model and compare outputs
755745
n, c, h, w = input_blob_dims
756746
data = np.random.randn(n, c, h, w).astype(np.float32)
757747
inputs = [data]
758748
_, c2_outputs = c2_native_run_net(c2_init_net, c2_predict_net, inputs)
759749
del _
760750

761-
with open(os.path.join(model_dir, 'value_info.json'), 'r') as value_info_conf:
762-
model = c2_onnx.caffe2_net_to_onnx_model(
763-
predict_net=c2_predict_net,
764-
init_net=c2_init_net,
765-
value_info=json.load(value_info_conf))
751+
model = c2_onnx.caffe2_net_to_onnx_model(
752+
predict_net=c2_predict_net,
753+
init_net=c2_init_net,
754+
value_info=value_info,
755+
)
766756
c2_ir = c2.prepare(model)
767757
onnx_outputs = c2_ir.run(inputs)
768758
self.assertSameOutputs(c2_outputs, onnx_outputs, decimal=decimal)

‎caffe2/python/onnx/tests/test_utils.py‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,6 @@
1010
import unittest
1111

1212
import numpy as np
13-
from caffe2.python.models.download import downloadFromURLToFile, getURLFromName, deleteDirectory
14-
1513

1614
class TestCase(unittest.TestCase):
1715

@@ -31,27 +29,3 @@ def add_test_case(self, name, test_func):
3129
raise ValueError('Duplicated test name: {}'.format(name))
3230
setattr(self, name, test_func)
3331

34-
35-
class DownloadingTestCase(TestCase):
36-
37-
def _download(self, model):
38-
model_dir = self._model_dir(model)
39-
assert not os.path.exists(model_dir)
40-
os.makedirs(model_dir)
41-
for f in ['predict_net.pb', 'init_net.pb', 'value_info.json']:
42-
url = getURLFromName(model, f)
43-
dest = os.path.join(model_dir, f)
44-
try:
45-
try:
46-
downloadFromURLToFile(url, dest,
47-
show_progress=False)
48-
except TypeError:
49-
# show_progress not supported prior to
50-
# Caffe2 78c014e752a374d905ecfb465d44fa16e02a28f1
51-
# (Sep 17, 2017)
52-
downloadFromURLToFile(url, dest)
53-
except Exception as e:
54-
print("Abort: {reason}".format(reason=e))
55-
print("Cleaning up...")
56-
deleteDirectory(model_dir)
57-
raise AssertionError("Test model downloading failed")

‎caffe2/python/trt/test_trt.py‎

Lines changed: 10 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,11 @@
99
import onnx.defs
1010
from onnx.helper import make_node, make_graph, make_tensor, make_tensor_value_info, make_model
1111
from onnx.backend.base import namedtupledict
12-
from caffe2.python.models.download import downloadFromURLToFile, getURLFromName, deleteDirectory
12+
from caffe2.python.models.download import ModelDownloader
1313
import caffe2.python.onnx.backend as c2
1414
from caffe2.python.onnx.workspace import Workspace
1515
from caffe2.python.trt.transform import convert_onnx_model_to_trt_op, transform_caffe2_net
16-
from caffe2.python.onnx.tests.test_utils import TestCase, DownloadingTestCase
16+
from caffe2.python.onnx.tests.test_utils import TestCase
1717
import numpy as np
1818
import os.path
1919
import json
@@ -80,6 +80,9 @@ def _download_onnx_model(model_name, opset_version):
8080
return model_dir
8181

8282
class TensorRTOpTest(TestCase):
83+
def setUp(self):
84+
self.opset_version = onnx.defs.onnx_opset_version()
85+
8386
def _test_relu_graph(self, X, batch_size, trt_max_batch_size):
8487
node_def = make_node("Relu", ["X"], ["Y"])
8588
Y_c2 = c2.run_node(node_def, {"X": X})
@@ -114,8 +117,7 @@ def test_relu_graph_big_batch(self):
114117
X = np.random.randn(52, 1, 3, 2).astype(np.float32)
115118
self._test_relu_graph(X, 52, 50)
116119

117-
def _test_onnx_importer(self, model_name, data_input_index,
118-
opset_version = onnx.defs.onnx_opset_version()):
120+
def _test_onnx_importer(self, model_name, data_input_index, opset_version=onnx.defs.onnx_opset_version()):
119121
model_dir = _download_onnx_model(model_name, opset_version)
120122
model_def = onnx.load(os.path.join(model_dir, 'model.onnx'))
121123
input_blob_dims = [int(x.dim_value) for x in model_def.graph.input[data_input_index].type.tensor_type.shape.dim]
@@ -176,31 +178,9 @@ def test_vgg19(self):
176178
self._test_onnx_importer('vgg19', -2, 9)
177179

178180

179-
class TensorRTTransformTest(DownloadingTestCase):
180-
def _model_dir(self, model):
181-
caffe2_home = os.path.expanduser(os.getenv('CAFFE2_HOME', '~/.caffe2'))
182-
models_dir = os.getenv('CAFFE2_MODELS', os.path.join(caffe2_home, 'models'))
183-
return os.path.join(models_dir, model)
184-
185-
def _get_c2_model(self, model_name):
186-
model_dir = self._model_dir(model_name)
187-
if not os.path.exists(model_dir):
188-
self._download(model_name)
189-
c2_predict_pb = os.path.join(model_dir, 'predict_net.pb')
190-
c2_predict_net = caffe2_pb2.NetDef()
191-
with open(c2_predict_pb, 'rb') as f:
192-
c2_predict_net.ParseFromString(f.read())
193-
c2_predict_net.name = model_name
194-
195-
c2_init_pb = os.path.join(model_dir, 'init_net.pb')
196-
c2_init_net = caffe2_pb2.NetDef()
197-
with open(c2_init_pb, 'rb') as f:
198-
c2_init_net.ParseFromString(f.read())
199-
c2_init_net.name = model_name + '_init'
200-
201-
with open(os.path.join(model_dir, 'value_info.json')) as f:
202-
value_info = json.load(f)
203-
return c2_init_net, c2_predict_net, value_info
181+
class TensorRTTransformTest(TestCase):
182+
def setUp(self):
183+
self.model_downloader = ModelDownloader()
204184

205185
def _add_head_tail(self, pred_net, new_head, new_tail):
206186
orig_head = pred_net.external_input[0]
@@ -226,14 +206,13 @@ def _add_head_tail(self, pred_net, new_head, new_tail):
226206
pred_net.op.extend([tail])
227207
pred_net.external_output[0] = new_tail
228208

229-
230209
@unittest.skipIf(not workspace.C.use_trt, "No TensortRT support")
231210
def test_resnet50_core(self):
232211
N = 2
233212
warmup = 20
234213
repeat = 100
235214
print("Batch size: {}, repeat inference {} times, warmup {} times".format(N, repeat, warmup))
236-
init_net, pred_net, _ = self._get_c2_model('resnet50')
215+
init_net, pred_net, _ = self.model_downloader.get_c2_model('resnet50')
237216
self._add_head_tail(pred_net, 'real_data', 'real_softmax')
238217
input_blob_dims = (N, 3, 224, 224)
239218
input_name = "real_data"

0 commit comments

Comments
 (0)