Add SuperAttention

lib/xlayers/super_module.py

@@ -1,13 +1,18 @@
 #####################################################
-# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.01 #
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
 #####################################################
 
 import abc
+from typing import Optional, Union, Callable
+import torch
 import torch.nn as nn
 from enum import Enum
 
 import spaces
 
+IntSpaceType = Union[int, spaces.Integer, spaces.Categorical]
+BoolSpaceType = Union[bool, spaces.Categorical]
+
 
 class SuperRunMode(Enum):
     """This class defines the enumerations for Super Model Running Mode."""
@@ -24,6 +29,7 @@ class SuperModule(abc.ABC, nn.Module):
         super(SuperModule, self).__init__()
         self._super_run_type = SuperRunMode.Default
         self._abstract_child = None
+        self._verbose = False
 
     def set_super_run_type(self, super_run_type):
         def _reset_super_run(m):
@@ -32,6 +38,13 @@ class SuperModule(abc.ABC, nn.Module):
 
         self.apply(_reset_super_run)
 
+    def apply_verbose(self, verbose):
+        def _reset_verbose(m):
+            if isinstance(m, SuperModule):
+                m._verbose = verbose
+
+        self.apply(_reset_verbose)
+
     def apply_candidate(self, abstract_child):
         if not isinstance(abstract_child, spaces.VirtualNode):
             raise ValueError(
@@ -51,6 +64,10 @@ class SuperModule(abc.ABC, nn.Module):
     def abstract_child(self):
         return self._abstract_child
 
+    @property
+    def verbose(self):
+        return self._verbose
+
     @abc.abstractmethod
     def forward_raw(self, *inputs):
         """Use the largest candidate for forward. Similar to the original PyTorch model."""
@@ -60,12 +77,41 @@ class SuperModule(abc.ABC, nn.Module):
     def forward_candidate(self, *inputs):
         raise NotImplementedError
 
+    @property
+    def name_with_id(self):
+        return "name={:}, id={:}".format(self.__class__.__name__, id(self))
+
+    def get_shape_str(self, tensors):
+        if isinstance(tensors, (list, tuple)):
+            shapes = [self.get_shape_str(tensor) for tensor in tensors]
+            if len(shapes) == 1:
+                return shapes[0]
+            else:
+                return ", ".join(shapes)
+        elif isinstance(tensors, (torch.Tensor, nn.Parameter)):
+            return str(tuple(tensors.shape))
+        else:
+            raise TypeError("Invalid input type: {:}.".format(type(tensors)))
+
     def forward(self, *inputs):
+        if self.verbose:
+            print(
+                "[{:}] inputs shape: {:}".format(
+                    self.name_with_id, self.get_shape_str(inputs)
+                )
+            )
         if self.super_run_type == SuperRunMode.FullModel:
-            return self.forward_raw(*inputs)
+            outputs = self.forward_raw(*inputs)
         elif self.super_run_type == SuperRunMode.Candidate:
-            return self.forward_candidate(*inputs)
+            outputs = self.forward_candidate(*inputs)
         else:
             raise ModeError(
                 "Unknown Super Model Run Mode: {:}".format(self.super_run_type)
             )
+        if self.verbose:
+            print(
+                "[{:}] outputs shape: {:}".format(
+                    self.name_with_id, self.get_shape_str(outputs)
+                )
+            )
+        return outputs