Merge pull request #70 from NVIDIA-AI-IOT/mobilenet_v2

jaybdub · web-flow · commit b442282c92d1 · 2019-09-25T16:26:59.000-04:00
added mobilenet_v2 to module tests
diff --git a/test.sh b/test.sh
@@ -27,3 +27,4 @@ python3 -m torch2trt.test -o $OUTPUT_FILE --name vgg11_bn
 python3 -m torch2trt.test -o $OUTPUT_FILE --name vgg13_bn
 python3 -m torch2trt.test -o $OUTPUT_FILE --name vgg16_bn
 python3 -m torch2trt.test -o $OUTPUT_FILE --name vgg19_bn
+python3 -m torch2trt.test -o $OUTPUT_FILE --name mobilenet_v2
diff --git a/torch2trt/module_test.py b/torch2trt/module_test.py
@@ -35,6 +35,7 @@ def module_name(self):
     ModuleTest(torchvision.models.vgg13_bn, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),
     ModuleTest(torchvision.models.vgg16_bn, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),
     ModuleTest(torchvision.models.vgg19_bn, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),
+    ModuleTest(torchvision.models.mobilenet_v2, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),
 ]
 
 

Original file line number	Diff line number	Diff line change
`@@ -35,6 +35,7 @@ def module_name(self):`
`35`	`35`	`ModuleTest(torchvision.models.vgg13_bn, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),`
`36`	`36`	`ModuleTest(torchvision.models.vgg16_bn, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),`
`37`	`37`	`ModuleTest(torchvision.models.vgg19_bn, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),`
	`38`	`+ ModuleTest(torchvision.models.mobilenet_v2, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),`
`38`	`39`	`]`
`39`	`40`
`40`	`41`