class

CompressionAware

extendsModule

CompressionAware(model: Module, weights: WeightPrecision | Palettize | Sparsify)

source edit

A model that trains against the compression it will be exported with.

Wrap a model, fine-tune the wrapper as if it were the model, then settle and export with the same weights argument. The forward pass uses compressed weights throughout, so the loss the optimizer sees is the loss the package will have.

Parameters

modelnn.Module

Model to fine-tune. It is used in place rather than copied — its parameters are what training updates — so keep a checkpoint if the uncompressed weights are still wanted.

weightsWeightPrecision or Palettize or Sparsify

The same value the export will be given. Anything else trains the model for a constraint the package will not have.

Attributes

coveredlist of str

Qualified names of the weights being compressed. Rank-one parameters and anything under the export's size threshold are absent, because the export leaves those in floating point too.

Notes

refit is the caller's to schedule. The codebook is fitted once at construction and the assignment onto it is recomputed every forward, so a model drifts away from its palette over a long fine-tune; refitting every few hundred steps follows it back. Refitting every step is not better — it moves the target the model is chasing.

Examples

>>> import shutil, tempfile
>>> import lucid, lucid.nn as nn, lucid.nn.functional as F
>>> import lucid.coreml as cml
>>> model = nn.Sequential(
...     nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(),
...     nn.Conv2d(32, 32, 3, padding=1), nn.ReLU(),
...     nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32, 10),
... )
>>> loader = [
...     (lucid.randn(4, 3, 16, 16), lucid.randint(0, 10, (4,))) for _ in range(3)
... ]
>>> aware = cml.CompressionAware(model, weights=cml.Palettize(bits=4))
>>> aware.covered               # the one weight big enough to compress
['2.weight']
>>> optimizer = lucid.optim.SGD(aware.parameters(), lr=0.01)
>>> for x, y in loader:
...     loss = F.cross_entropy(aware(x), y)
...     loss.backward()
...     optimizer.step()
...     optimizer.zero_grad()
>>> aware.refit()               # follow the weights as they move
>>> settled = aware.settle().eval()
>>> x, room = lucid.randn(1, 3, 16, 16), tempfile.mkdtemp()
>>> with cml.export(settled, x, f"{room}/aware.mlpackage",
...                 weights=cml.Palettize(bits=4)) as package:
...     print(package.verify(settled, x, relative=True) < 1e-5)
True
>>> shutil.rmtree(room)

Used by 1

lucid.coreml

Constructors

dunder

init

→None

__init__(model: Module, weights: WeightPrecision | Palettize | Sparsify)

source edit

Properties

prop

covered

→list[str]

covered: list[str]

source edit

Weights this is training against, in the model's own names.

Instance methods

forward

→Tensor | tuple[Tensor, ...]

forward(args: Tensor = (), kwargs: object = {})

source edit

Run the model with every covered weight compressed.

refit

→None

refit()

source edit

Re-fit every codebook, scale or mask to the current weights.

settle

→nn.Module

settle()

source edit

Write the compressed values into the model and hand it back.

Until this is called the parameters are still full precision and only the forward pass sees the compression; exporting then would fit a palette to weights that sit near one rather than on it.

Returns

nn.Module

The model that was passed in, with every covered weight replaced by the value the compression gives it.

>>> import shutil, tempfile >>> import lucid, lucid.nn as nn, lucid.nn.functional as F >>> import lucid.coreml as cml >>> model = nn.Sequential( ... nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), ... nn.Conv2d(32, 32, 3, padding=1), nn.ReLU(), ... nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32, 10), ... ) >>> loader = [ ... (lucid.randn(4, 3, 16, 16), lucid.randint(0, 10, (4,))) for _ in range(3) ... ] >>> aware = cml.CompressionAware(model, weights=cml.Palettize(bits=4)) >>> aware.covered # the one weight big enough to compress ['2.weight'] >>> optimizer = lucid.optim.SGD(aware.parameters(), lr=0.01) >>> for x, y in loader: ... loss = F.cross_entropy(aware(x), y) ... loss.backward() ... optimizer.step() ... optimizer.zero_grad() >>> aware.refit() # follow the weights as they move >>> settled = aware.settle().eval() >>> x, room = lucid.randn(1, 3, 16, 16), tempfile.mkdtemp() >>> with cml.export(settled, x, f"{room}/aware.mlpackage", ... weights=cml.Palettize(bits=4)) as package: ... print(package.verify(settled, x, relative=True) < 1e-5) True >>> shutil.rmtree(room)