8
8
import torch
9
9
10
10
from detectron2 import model_zoo
11
+ from detectron2 .export import Caffe2Model , Caffe2Tracer , add_export_config
11
12
from detectron2 .utils .logger import setup_logger
12
13
from detectron2 .utils .testing import get_sample_coco_image
13
14
14
15
15
- @unittest .skipIf (os .environ .get ("CI" ), "Require COCO data and model zoo." )
16
+ # TODO: this test requires manifold access, see: T88318502
17
+ # Running it on CircleCI causes crash, not sure why.
18
+ @unittest .skipIf (os .environ .get ("CIRCLECI" ), "Caffe2 tests crash on CircleCI." )
16
19
class TestCaffe2Export (unittest .TestCase ):
17
20
def setUp (self ):
18
21
setup_logger ()
19
22
20
23
def _test_model (self , config_path , device = "cpu" ):
21
- # requires extra dependencies
22
- from detectron2 .export import Caffe2Model , add_export_config , Caffe2Tracer
23
-
24
24
cfg = model_zoo .get_config (config_path )
25
25
add_export_config (cfg )
26
26
cfg .MODEL .DEVICE = device
@@ -29,27 +29,25 @@ def _test_model(self, config_path, device="cpu"):
29
29
inputs = [{"image" : get_sample_coco_image ()}]
30
30
tracer = Caffe2Tracer (cfg , model , copy .deepcopy (inputs ))
31
31
32
- c2_model = tracer .export_caffe2 ()
33
-
34
32
with tempfile .TemporaryDirectory (prefix = "detectron2_unittest" ) as d :
35
- c2_model .save_protobuf (d )
36
- c2_model .save_graph (os .path .join (d , "test.svg" ), inputs = copy .deepcopy (inputs ))
33
+ if not os .environ .get ("CI" ):
34
+ # This requires onnx, which is not yet available on public CI
35
+ c2_model = tracer .export_caffe2 ()
36
+ c2_model .save_protobuf (d )
37
+ c2_model .save_graph (os .path .join (d , "test.svg" ), inputs = copy .deepcopy (inputs ))
37
38
38
- c2_model = Caffe2Model .load_protobuf (d )
39
- c2_model (inputs )[0 ]["instances" ]
39
+ c2_model = Caffe2Model .load_protobuf (d )
40
+ c2_model (inputs )[0 ]["instances" ]
40
41
41
42
ts_model = tracer .export_torchscript ()
42
43
ts_model .save (os .path .join (d , "model.ts" ))
43
44
44
45
def testMaskRCNN (self ):
45
- # TODO: this test requires manifold access, see: T88318502
46
46
self ._test_model ("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml" )
47
47
48
48
@unittest .skipIf (not torch .cuda .is_available (), "CUDA not available" )
49
49
def testMaskRCNNGPU (self ):
50
- # TODO: this test requires manifold access, see: T88318502
51
50
self ._test_model ("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml" , device = "cuda" )
52
51
53
52
def testRetinaNet (self ):
54
- # TODO: this test requires manifold access, see: T88318502
55
53
self ._test_model ("COCO-Detection/retinanet_R_50_FPN_3x.yaml" )
0 commit comments