What to do about a model that draws random numbers in forward.
Core ML has no random operation. A draw whose inputs are all
constants — which every randn is — folds at build time, so a
package written the obvious way returns one fixed sample for the life
of the file. A variational encoder exported that way has a latent
that never moves, and nothing about the package looks wrong.
REFUSED is the default and says so rather than shipping that.
AS_INPUT lifts each draw to an input the caller fills. The
reparameterisation mu + sigma * eps becomes a function of eps
instead of a function that makes one, which is what the model always
was — the draw was never part of the network. Seven families in the
zoo are exportable only this way: variational encoders, world models
with a stochastic latent, and score-based samplers.
The handle draws for the caller who does not, so a package still answers differently on each prediction the way the eager model does.
Examples
>>> import shutil, tempfile
>>> import lucid, lucid.nn as nn, lucid.coreml as cml
>>> class Encoder(nn.Module):
... def __init__(self):
... super().__init__()
... self.mu = nn.Linear(8, 4)
... self.logvar = nn.Linear(8, 4)
... def forward(self, x):
... mu, logvar = self.mu(x), self.logvar(x)
... return mu + (logvar * 0.5).exp() * lucid.randn(1, 4)
>>> vae, x, room = Encoder().eval(), lucid.randn(1, 8), tempfile.mkdtemp()
>>> package = cml.export(vae, x, f"{room}/vae.mlpackage",
... draws=cml.Draws.AS_INPUT)
>>> package.noise_inputs
[('noise_0', (1, 4))]
Pass nothing for it and the handle draws, so the package samples the
way the model does:
>>> float((package.predict(x) - package.predict(x)).abs().max()) > 0
True
Pass one and the package is a deterministic function of it:
>>> eps = lucid.randn(1, 4)
>>> a = package.predict({"input": x, "noise_0": eps})
>>> b = package.predict({"input": x, "noise_0": eps})
>>> float((a - b).abs().max())
0.0
>>> package.close()
>>> shutil.rmtree(room)