Use forward() instead of __call__()

Author: koki0702

dezero/layers.py

@@ -1,4 +1,5 @@
 import os
+import weakref
 import numpy as np
 import dezero.functions as F
 from dezero import cuda
@@ -18,11 +19,16 @@ def __setattr__(self, name, value):
             self._params.add(name)
         super().__setattr__(name, value)
 
-    def __call__(self, *args, **kwargs):
-        return self.forward(*args, **kwargs)
+    def __call__(self, *inputs):
+        outputs = self.forward(*inputs)
+        if not isinstance(outputs, tuple):
+            outputs = (outputs,)
+        self.inputs = [weakref.ref(x) for x in inputs]
+        self.outputs = [weakref.ref(y) for y in outputs]
+        return outputs if len(outputs) > 1 else outputs[0]
 
-    def forward(self, *args, **kwargs):
-        return self.__call__(*args, **kwargs)
+    def forward(self, inputs):
+        raise NotImplementedError()
 
     def params(self):
         for name in self._params:
@@ -101,7 +107,7 @@ def _init_W(self, xp=np):
         W_data = xp.random.randn(I, O).astype(self.dtype) * np.sqrt(1 / I)
         self.W.data = W_data
 
-    def __call__(self, x):
+    def forward(self, x):
         if self.W.data is None:
             self.in_size = x.shape[1]
             xp = cuda.get_array_module(x)
@@ -150,7 +156,7 @@ def _init_W(self, xp=np):
         W_data = xp.random.randn(OC, C, KH, KW).astype(self.dtype) * scale
         self.W.data = W_data
 
-    def __call__(self, x):
+    def forward(self, x):
         if self.W.data is None:
             self.in_channels = x.shape[1]
             xp = cuda.get_array_module(x)
@@ -199,7 +205,7 @@ def _init_W(self, xp=np):
         W_data = xp.random.randn(C, OC, KH, KW).astype(self.dtype) * scale
         self.W.data = W_data
 
-    def __call__(self, x):
+    def forward(self, x):
         if self.W.data is None:
             self.in_channels = x.shape[1]
             xp = cuda.get_array_module(x)
@@ -231,7 +237,7 @@ def __init__(self, hidden_size, in_size=None):
     def reset_state(self):
         self.h = None
 
-    def __call__(self, x):
+    def forward(self, x):
         if self.h is None:
             h_new = F.tanh(self.x2h(x))
         else:
@@ -259,7 +265,7 @@ def reset_state(self):
         self.h = None
         self.c = None
 
-    def __call__(self, x):
+    def forward(self, x):
         if self.h is None:
             f = F.sigmoid(self.x2f(x))
             i = F.sigmoid(self.x2i(x))

dezero/models.py

@@ -11,7 +11,7 @@
 # =============================================================================
 class Model(Layer):
     def plot(self, *inputs, to_file='model.png'):
-        y = self.__call__(*inputs)
+        y = self.forward(*inputs)
         return utils.plot_dot_graph(y, verbose=True, to_file=to_file)
 
 
@@ -23,7 +23,7 @@ def __init__(self, *layers):
             setattr(self, 'l' + str(i), layer)
             self.layers.append(layer)
 
-    def __call__(self, x):
+    def forward(self, x):
         for layer in self.layers:
             x = layer(x)
         return x
@@ -40,7 +40,7 @@ def __init__(self, fc_output_sizes, activation=F.sigmoid):
             setattr(self, 'l' + str(i), layer)
             self.layers.append(layer)
 
-    def __call__(self, x):
+    def forward(self, x):
         for l in self.layers[:-1]:
             x = self.activation(l(x))
         return self.layers[-1](x)
@@ -75,7 +75,7 @@ def __init__(self, pretrained=False):
             weights_path = utils.get_file(VGG16.WEIGHTS_PATH)
             self.load_weights(weights_path)
 
-    def __call__(self, x):
+    def forward(self, x):
         x = F.relu(self.conv1_1(x))
         x = F.relu(self.conv1_2(x))
         x = F.pooling(x, 2, 2)
@@ -143,7 +143,7 @@ def __init__(self, n_layers=152, pretrained=False):
             weights_path = utils.get_file(ResNet.WEIGHTS_PATH.format(n_layers))
             self.load_weights(weights_path)
 
-    def __call__(self, x):
+    def forward(self, x):
         x = F.relu(self.bn1(self.conv1(x)))
         x = F.pooling(x, kernel_size=3, stride=2)
         x = self.res2(x)
@@ -271,5 +271,5 @@ class SqueezeNet(Model):
     def __init__(self, pretrained=False):
         pass
 
-    def __call__(self, x):
+    def forward(self, x):
         pass

dezero/utils.py

@@ -236,15 +236,15 @@ def numerical_grad(f, x, *args, **kwargs):
 
     Args:
         f (callable): A function which gets `Variable`s and returns `Variable`s.
-        x (`ndarray` or `dezero.Variable`): A traget `Variable` for computing
+        x (`ndarray` or `dezero.Variable`): A target `Variable` for computing
             the gradient.
         *args: If `f` needs variables except `x`, you can specify with this
             argument.
         **kwargs: If `f` needs keyword variables, you can specify with this
             argument.
 
     Returns:
-        `ndarray`:
+        `ndarray`: Gradient.
     """
     eps = 1e-4
 

examples/spiral.py

@@ -25,7 +25,7 @@ def __init__(self, hidden_size, out_size):
         self.l2 = L.Linear(out_size)
         self.bn1 = L.BatchNorm()
 
-    def __call__(self, x):
+    def forward(self, x):
         y = F.sigmoid(self.bn1(self.l1(x)))
         y = self.l2(y)
         return y

examples/vae.py

@@ -27,7 +27,7 @@ def __init__(self, latent_size):
         self.linear2 = L.Linear(latent_size)
         self.linear3 = L.Linear(latent_size)
 
-    def __call__(self, x):
+    def forward(self, x):
         x = F.relu(self.conv1(x))
         x = F.relu(self.conv2(x))
         x = F.relu(self.conv3(x))
@@ -53,7 +53,7 @@ def __init__(self):
         self.deconv = L.Deconv2d(32, kernel_size=4, stride=2, pad=1)
         self.conv = L.Conv2d(1, kernel_size=3, stride=1, pad=1)
 
-    def __call__(self, x):
+    def forward(self, x):
         x = F.relu(self.linear(x))
         x = F.reshape(x, (-1,) + self.to_shape)  # reshape to (-1, C, H, W)
         x = F.relu(self.deconv(x))
@@ -68,7 +68,7 @@ def __init__(self, latent_size):
         self.encoder = Encoder(latent_size)
         self.decoder = Decoder()
 
-    def __call__(self, x, C=1.0, k=1):
+    def forward(self, x, C=1.0, k=1):
         """Call loss function of VAE.
         The loss value is equal to ELBO (Evidence Lower Bound)
         multiplied by -1.

steps/step45.py

@@ -23,7 +23,7 @@ def __init__(self, hidden_size, out_size):
         self.l1 = L.Linear(hidden_size)
         self.l2 = L.Linear(out_size)
 
-    def __call__(self, x):
+    def forward(self, x):
         y = F.sigmoid(self.l1(x))
         y = self.l2(y)
         return y