3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
6 # Written by Francois Fleuret <francois@fleuret.org>
8 from torch import nn, Tensor
10 ######################################################################
13 class LazyLinear(nn.Module):
14 def __init__(self, out_dim, bias=True):
16 self.out_dim = out_dim
21 x = x.view(x.size(0), -1)
25 self.core = nn.Linear(x.size(1), self.out_dim, self.bias)
27 raise RuntimeError("Undefined LazyLinear core in inference mode.")
31 def named_parameters(self, memo=None, prefix=""):
32 assert self.core is not None, "Parameters not yet defined"
33 return super().named_parameters(memo, prefix)
36 ######################################################################
38 if __name__ == "__main__":
39 model = nn.Sequential(
40 nn.Conv2d(3, 8, kernel_size=5),
41 nn.ReLU(inplace=True),
43 nn.ReLU(inplace=True),
49 input = Tensor(100, 3, 32, 32).normal_()
53 for n, x in model.named_parameters():