Add SuperTransformer

lib/xlayers/super_attention.py

@@ -37,10 +37,7 @@ class SuperAttention(SuperModule):
         self._proj_dim = proj_dim
         self._num_heads = num_heads
         self._qkv_bias = qkv_bias
-        # head_dim = dim // num_heads
-        # self.scale = qk_scale or math.sqrt(head_dim)
 
-        # self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
         self.q_fc = SuperLinear(input_dim, input_dim, bias=qkv_bias)
         self.k_fc = SuperLinear(input_dim, input_dim, bias=qkv_bias)
         self.v_fc = SuperLinear(input_dim, input_dim, bias=qkv_bias)

lib/xlayers/super_core.py

@@ -2,6 +2,8 @@
 # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
 #####################################################
 from .super_module import SuperRunMode
+from .super_module import IntSpaceType
+
 from .super_module import SuperModule
 from .super_container import SuperSequential
 from .super_linear import SuperLinear
@@ -9,3 +11,6 @@ from .super_linear import SuperMLPv1, SuperMLPv2
 from .super_norm import SuperLayerNorm1D
 from .super_attention import SuperAttention
 from .super_transformer import SuperTransformerEncoderLayer
+
+from .super_trade_stem import SuperAlphaEBDv1
+from .super_positional_embedding import SuperPositionalEncoder

lib/xlayers/super_linear.py

@@ -109,7 +109,7 @@ class SuperLinear(SuperModule):
 
     def extra_repr(self) -> str:
         return "in_features={:}, out_features={:}, bias={:}".format(
-            self.in_features, self.out_features, self.bias
+            self._in_features, self._out_features, self._bias
         )
 
 

lib/xlayers/super_norm.py

@@ -75,8 +75,10 @@ class SuperLayerNorm1D(SuperModule):
         return F.layer_norm(input, (self.in_dim,), self.weight, self.bias, self.eps)
 
     def extra_repr(self) -> str:
-        return "{in_dim}, eps={eps}, " "elementwise_affine={elementwise_affine}".format(
-            in_dim=self._in_dim,
-            eps=self._eps,
-            elementwise_affine=self._elementwise_affine,
+        return (
+            "shape={in_dim}, eps={eps}, elementwise_affine={elementwise_affine}".format(
+                in_dim=self._in_dim,
+                eps=self._eps,
+                elementwise_affine=self._elementwise_affine,
+            )
         )

lib/xlayers/super_positional_embedding.py

@@ -0,0 +1,68 @@
+#####################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.02 #
+#####################################################
+import torch
+import torch.nn as nn
+import math
+
+import spaces
+from .super_module import SuperModule
+from .super_module import IntSpaceType
+
+
+class SuperPositionalEncoder(SuperModule):
+    """Attention Is All You Need: https://arxiv.org/pdf/1706.03762.pdf
+    https://github.com/pytorch/examples/blob/master/word_language_model/model.py#L65
+    """
+
+    def __init__(self, d_model: IntSpaceType, max_seq_len: int, dropout: float = 0.1):
+        super(SuperPositionalEncoder, self).__init__()
+        self._d_model = d_model
+        # create constant 'pe' matrix with values dependant on
+        # pos and i
+        self.dropout = nn.Dropout(p=dropout)
+        self.register_buffer("pe", self.create_pos_embed(max_seq_len, self.d_model))
+
+    @property
+    def d_model(self):
+        return spaces.get_max(self._d_model)
+
+    @property
+    def abstract_search_space(self):
+        root_node = spaces.VirtualNode(id(self))
+        if not spaces.is_determined(self._d_model):
+            root_node.append("_d_model", self._d_model.abstract(reuse_last=True))
+        return root_node
+
+    def create_pos_embed(self, max_seq_len, d_model):
+        pe = torch.zeros(max_seq_len, d_model)
+        for pos in range(max_seq_len):
+            for i in range(0, d_model):
+                div = 10000 ** ((i // 2) * 2 / d_model)
+                value = pos / div
+                if i % 2 == 0:
+                    pe[pos, i] = math.sin(value)
+                else:
+                    pe[pos, i] = math.cos(value)
+        return pe.unsqueeze(0)
+
+    def forward_candidate(self, input: torch.Tensor) -> torch.Tensor:
+        batch, seq, fdim = input.shape[:3]
+        embeddings = self.pe[:, :seq]
+        if not spaces.is_determined(self._d_model):
+            expected_d_model = self.abstract_child["_d_model"].value
+        else:
+            expected_d_model = spaces.get_determined_value(self._d_model)
+        assert fdim == expected_d_model, "{:} vs {:}".format(fdim, expected_d_model)
+
+        embeddings = torch.nn.functional.interpolate(
+            embeddings, size=(expected_d_model), mode="linear", align_corners=True
+        )
+        outs = self.dropout(input + embeddings)
+        return outs
+
+    def forward_raw(self, input: torch.Tensor) -> torch.Tensor:
+        batch, seq, fdim = input.shape[:3]
+        embeddings = self.pe[:, :seq]
+        outs = self.dropout(input + embeddings)
+        return outs

lib/xlayers/super_trade_stem.py

@@ -0,0 +1,63 @@
+#####################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
+#####################################################
+from __future__ import division
+from __future__ import print_function
+
+import math
+from functools import partial
+from typing import Optional, Text
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import spaces
+from .super_linear import SuperLinear
+from .super_module import SuperModule
+from .super_module import IntSpaceType
+
+
+class SuperAlphaEBDv1(SuperModule):
+    """A simple layer to convert the raw trading data from 1-D to 2-D data and apply an FC layer."""
+
+    def __init__(self, d_feat: int, embed_dim: IntSpaceType):
+        super(SuperAlphaEBDv1, self).__init__()
+        self._d_feat = d_feat
+        self._embed_dim = embed_dim
+        self.proj = SuperLinear(d_feat, embed_dim)
+
+    @property
+    def embed_dim(self):
+        return spaces.get_max(self._embed_dim)
+
+    @property
+    def abstract_search_space(self):
+        root_node = spaces.VirtualNode(id(self))
+        space = self.proj.abstract_search_space
+        if not spaces.is_determined(space):
+            root_node.append("proj", space)
+        if not spaces.is_determined(self._embed_dim):
+            root_node.append("_embed_dim", self._embed_dim.abstract(reuse_last=True))
+        return root_node
+
+    def apply_candidate(self, abstract_child: spaces.VirtualNode):
+        super(SuperAlphaEBDv1, self).apply_candidate(abstract_child)
+        if "proj" in abstract_child:
+            self.proj.apply_candidate(abstract_child["proj"])
+
+    def forward_candidate(self, input: torch.Tensor) -> torch.Tensor:
+        x = input.reshape(len(input), self._d_feat, -1)  # [N, F*T] -> [N, F, T]
+        x = x.permute(0, 2, 1)  # [N, F, T] -> [N, T, F]
+        if not spaces.is_determined(self._embed_dim):
+            embed_dim = self.abstract_child["_embed_dim"].value
+        else:
+            embed_dim = spaces.get_determined_value(self._embed_dim)
+        out = self.proj(x) * math.sqrt(embed_dim)
+        return out
+
+    def forward_raw(self, input: torch.Tensor) -> torch.Tensor:
+        x = input.reshape(len(input), self._d_feat, -1)  # [N, F*T] -> [N, F, T]
+        x = x.permute(0, 2, 1)  # [N, F, T] -> [N, T, F]
+        out = self.proj(x) * math.sqrt(self.embed_dim)
+        return out