From: Francois Fleuret Date: Thu, 10 May 2018 09:19:45 +0000 (+0200) Subject: Initial commit. X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=commitdiff_plain;h=7443d768bc437889659ba3ed737297f90fe1922e;p=pytorch.git Initial commit. --- 7443d768bc437889659ba3ed737297f90fe1922e diff --git a/lazy_linear.py b/lazy_linear.py new file mode 100755 index 0000000..7c9e398 --- /dev/null +++ b/lazy_linear.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python-for-pytorch + +from torch import nn, Tensor + +########## + +class LazyLinear(nn.Module): + + def __init__(self, out_dim, bias = True): + super(LazyLinear, self).__init__() + self.out_dim = out_dim + self.bias = bias + self.core = None + + def forward(self, x): + x = x.view(x.size(0), -1) + + if self.core is None: + if self.training: + self.core = nn.Linear(x.size(1), self.out_dim, self.bias) + else: + raise RuntimeError('Undefined LazyLinear core in inference mode.') + + return self.core(x) + +########## + +model = nn.Sequential(nn.Conv2d(1, 8, kernel_size = 5), + nn.ReLU(inplace = True), + LazyLinear(128), + nn.ReLU(inplace = True), + nn.Linear(128, 10)) + +# model.eval() + +input = Tensor(100, 1, 32, 32).normal_() + +output = model(input)