SqeezeNet meta data for create_cnn function

fastai/vision/learner.py

@@ -13,14 +13,20 @@
 def _default_split(m:nn.Module): return (m[1],)
 # Split a resnet style model
 def _resnet_split(m:nn.Module): return (m[0][6],m[1])
+# Split squeezenet model on maxpool layers
+def _squeezenet_split(m:nn.Module): return (m[0][0][5], m[0][0][8], m[1])
 
 _default_meta = {'cut':-1, 'split':_default_split}
 _resnet_meta  = {'cut':-2, 'split':_resnet_split }
+_squeezenet_meta = {'cut':-1, 'split': _squeezenet_split}
 
 model_meta = {
     models.resnet18 :{**_resnet_meta}, models.resnet34: {**_resnet_meta},
     models.resnet50 :{**_resnet_meta}, models.resnet101:{**_resnet_meta},
-    models.resnet152:{**_resnet_meta}}
+    models.resnet152:{**_resnet_meta},
+
+    models.squeezenet1_0:{**_squeezenet_meta},
+    models.squeezenet1_1:{**_squeezenet_meta}}
 
 def cnn_config(arch):
     "Get the metadata associated with `arch`."

fastai/vision/models/__init__.py

@@ -1,4 +1,5 @@
 from torchvision.models import ResNet,resnet18,resnet34,resnet50,resnet101,resnet152
+from torchvision.models import SqueezeNet,squeezenet1_0,squeezenet1_1
 from .darknet import *
 from .unet import *
 from .wrn import *

tests/test_vision_train.py

@@ -17,11 +17,22 @@ def no_bar():
     fastprogress.NO_BAR = False
 
 @pytest.fixture(scope="module")
-def learn():
+def mnist_tiny():
     path = untar_data(URLs.MNIST_TINY)
     data = ImageDataBunch.from_folder(path, ds_tfms=(rand_pad(2, 28), []), num_workers=2)
     data.normalize()
-    learn = Learner(data, simple_cnn((3,16,16,16,2), bn=True), metrics=[accuracy, error_rate])
+    return data
+
+@pytest.fixture(scope="module")
+def zero_image():
+    return Image(torch.zeros((3, 128, 128)))
+
+@pytest.fixture(scope="module")
+def learn(mnist_tiny):
+    # path = untar_data(URLs.MNIST_TINY)
+    # data = ImageDataBunch.from_folder(path, ds_tfms=(rand_pad(2, 28), []), num_workers=2)
+    # data.normalize()
+    learn = Learner(mnist_tiny, simple_cnn((3,16,16,16,2), bn=True), metrics=[accuracy, error_rate])
     learn.fit_one_cycle(3)
     return learn
 
@@ -107,3 +118,10 @@ def test_model_load_mem_leak(learn_large_unfit):
     used_after_reclaimed = gpu_mem_get_used()
     # XXX: not sure where 6MB get lost still but for now it's a small leak - need to test with a bigger model
     assert isclose(used_before, used_after_reclaimed, abs_tol=6),f"load() and used GPU RAM: before load(): {used_before}, after: {used_after}, after gc.collect() {used_after_reclaimed} used"
+
+@pytest.mark.slow
+@pytest.mark.parametrize('arch', [models.resnet18, models.squeezenet1_1])
+def test_models_meta(mnist_tiny, arch, zero_image):
+    learn = create_cnn(mnist_tiny, arch, metrics=[accuracy, error_rate])
+    pred = learn.predict(zero_image)
+    assert pred is not None