3 from torch import nn, Tensor
5 ######################################################################
7 class LazyLinear(nn.Module):
9 def __init__(self, out_dim, bias = True):
10 super(LazyLinear, self).__init__()
11 self.out_dim = out_dim
16 x = x.view(x.size(0), -1)
20 self.core = nn.Linear(x.size(1), self.out_dim, self.bias)
22 raise RuntimeError('Undefined LazyLinear core in inference mode.')
26 def named_parameters(self, memo=None, prefix=''):
27 assert self.core is not None, 'Parameters not yet defined'
28 return super(LazyLinear, self).named_parameters(memo, prefix)
30 ######################################################################
32 if __name__ == "__main__":
33 model = nn.Sequential(nn.Conv2d(3, 8, kernel_size = 5),
34 nn.ReLU(inplace = True),
36 nn.ReLU(inplace = True),
41 input = Tensor(100, 3, 32, 32).normal_()
45 for n, x in model.named_parameters():