version bump + fixes + hrank example

README.md

@@ -19,7 +19,7 @@ Mask and prune channels, using the default magnitude pruner.
 import torch, torchvision
 from upscale import MaskingManager, PruningManager
 
-x = torch.rand((1, 3, 224, 224)).cuda()
+x = torch.rand((1, 3, 224, 224), device='cuda')
 model = torchvision.models.get_model('resnet18', pretrained=True).cuda()  # get any pytorch model
 MaskingManager(model).importance().mask()
 PruningManager(model).compute([x]).prune()

examples/heuristic.py

@@ -0,0 +1,29 @@
+"""
+For licensing see accompanying LICENSE file.
+Copyright (C) 2023 Apple Inc. All Rights Reserved.
+
+How to customize heuristic used for pruning, in UPSCALE
+"""
+
+import torch, torchvision
+from upscale import MaskingManager, PruningManager
+from upscale.masking.importance import HRank, LAMP
+
+
+# for most heuristics, simply pass the heuristic to `.importance(...)`
+x = torch.rand((1, 3, 224, 224), device='cuda')
+model = torchvision.models.get_model('resnet18', pretrained=True).cuda()  # get any pytorch model
+MaskingManager(model).importance(LAMP()).mask()
+PruningManager(model).compute([x]).prune()
+
+# for HRank, we need to run several forward passes to collect feature map
+# statistics
+
+heuristic = HRank()
+heuristic.register(model)
+for _ in range(10):
+    model(torch.rand((1, 3, 224, 224), device='cuda'))
+heuristic.deregister(model)
+
+MaskingManager(model).importance(heuristic).mask()
+PruningManager(model).compute([x]).prune()

pyproject.toml

@@ -5,7 +5,7 @@
 
 [project]
 name = "apple-upscale"  # Required
-version = "0.1.0"  # Required
+version = "0.1.1"  # Required
 description = "Export utility for unconstrained channel pruned models"  # Optional
 readme = "README.md" # Optional
 requires-python = ">=3.7"

src/upscale/masking/importance.py

@@ -151,6 +151,7 @@ def hook(module, input, output):
                 logging.debug(f"Loaded precomputed rank for {module._name})")
 
         model._handles = []
+        model._name = ''
         is_hook_registered = False
         for name, module in model.named_modules():
             if isinstance(module, nn.Conv2d):

src/upscale/pruning/manager.py

@@ -32,7 +32,7 @@ class PruningManager(nn.Module):
 
     >>> manager = PruningManager(net)
     >>> pruned_inputs = [('layer1.1.conv1.weight', 5)]
-    >>> manager.compute([x], pruned_inputs=pruned_inputs)
+    >>> _ = manager.compute([x], pruned_inputs=pruned_inputs)
 
     Then, during training, you can use the mock-pruned model. This mock-pruned
     model applies masks instead of modifying the model itself.
@@ -43,7 +43,7 @@ class PruningManager(nn.Module):
     Finally, actually prune the model. Then run inference using the *original
     (now modified, in place) model.
 
-    >>> manager.prune()
+    >>> _ = manager.prune()
     >>> y_pruned = net(x)
 
     Check that both the mocked and pruned outputs match.

tests/test_api.py

@@ -0,0 +1,19 @@
+"""
+For licensing see accompanying LICENSE file.
+Copyright (C) 2023 Apple Inc. All Rights Reserved.
+
+Tests for the UPSCALE library external-facing API
+"""
+
+import torch
+import torch.nn as nn
+from torchvision.models import resnet18
+
+from upscale import MaskingManager, PruningManager
+
+
+def test_simple():
+    x = torch.rand((1, 3, 224, 224), device='cuda')
+    model = resnet18().cuda()
+    MaskingManager(model).importance().mask(amount=0.1)
+    PruningManager(model).compute([x]).prune()