Merge pull request #5 from kengz/scalartst

Support scalar output for transformers

.github/workflows/ci.yml

@@ -2,43 +2,49 @@ name: CI
 
 on:
   push:
-    branches: [ master ]
+    branches: [main]
   pull_request:
-    branches: [ master ]
+    branches: [main]
 
 jobs:
   build:
     runs-on: ubuntu-latest
+
     steps:
-    - uses: actions/checkout@v2
-    - name: Cache Conda
-      uses: actions/cache@v1
-      with:
-        path: /usr/share/miniconda/envs/torcharc
-        key: ${{ runner.os }}-conda-${{ hashFiles('environment.yml') }}
-        restore-keys: |
-          ${{ runner.os }}-conda-
-    - name: Setup Conda dependencies
-      uses: goanpeca/setup-miniconda@v1
-      with:
-         activate-environment: torcharc
-         environment-file: environment.yml
-         python-version: 3.8
-         auto-activate-base: false
-    - name: Conda info
-      shell: bash -l {0}
-      run: |
-        conda info
-        conda list
-    - name: Setup flake8 annotations
-      uses: rbialon/flake8-annotations@v1
-    - name: Lint with flake8
-      shell: bash -l {0}
-      run: |
-        pip install flake8
-        # exit-zero treats all errors as warnings.
-        flake8 . --ignore=E501 --count --exit-zero --statistics
-    - name: Run tests
-      shell: bash -l {0}
-      run: |
-        python setup.py test
+      - uses: actions/checkout@v2
+
+      - name: Cache Conda
+        uses: actions/cache@v2
+        with:
+          path: /usr/share/miniconda/envs/torcharc
+          key: ${{ runner.os }}-conda-${{ hashFiles('environment.yml') }}
+          restore-keys: |
+            ${{ runner.os }}-conda-
+
+      - name: Setup Conda dependencies
+        uses: conda-incubator/setup-miniconda@v2
+        with:
+          activate-environment: torcharc
+          environment-file: environment.yml
+          python-version: 3.8
+          auto-activate-base: false
+
+      - name: Conda info
+        shell: bash -l {0}
+        run: |
+          conda info
+          conda list
+
+      - name: Setup flake8 annotations
+        uses: rbialon/flake8-annotations@v1
+      - name: Lint with flake8
+        shell: bash -l {0}
+        run: |
+          pip install flake8
+          # exit-zero treats all errors as warnings.
+          flake8 . --ignore=E501 --count --exit-zero --statistics
+
+      - name: Run tests
+        shell: bash -l {0}
+        run: |
+          python setup.py test

setup.py

@@ -35,7 +35,7 @@ def run_tests(self):
 
 setup(
     name='torcharc',
-    version='0.0.5',
+    version='0.0.6',
     description='Build PyTorch networks by specifying architectures.',
     long_description='https://github.com/kengz/torcharc',
     keywords='torcharc',

torcharc/module/transformer/pytorch_tst.py

@@ -47,6 +47,7 @@ def __init__(
             q: int = 8,  # Dimension of queries and keys.
             v: int = 8,  # Dimension of values.
             chunk_mode: bool = 'chunk',
+            scalar_output: bool = False,
     ) -> None:
         super().__init__()
 
@@ -66,6 +67,7 @@ def __init__(
         self.decoders = nn.TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)
 
         self.out_linear = nn.Linear(d_model, out_channels)
+        self.scalar_output = scalar_output
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         '''
@@ -78,9 +80,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         encoding = self.encoders(encoding.transpose(0, 1))
 
         decoding = encoding
-        if self.pe is not None:  # position encoding
-            decoding = self.pe(decoding.transpose(0, 1)).transpose(0, 1)
-        decoding = self.decoders(decoding, encoding).transpose(0, 1)
+        if len(self.decoders.layers):
+            if self.pe is not None:  # position encoding
+                decoding = self.pe(decoding.transpose(0, 1)).transpose(0, 1)
+            decoding = self.decoders(decoding, encoding).transpose(0, 1)
+
+        if self.scalar_output:  # if want scalar instead of seq output, take the first index from seq
+            decoding = decoding[:, 0, :]
 
         output = self.out_linear(decoding)
         return output

torcharc/module/transformer/tst.py

@@ -49,6 +49,7 @@ def __init__(
             q: int = 8,  # Dimension of queries and keys.
             v: int = 8,  # Dimension of values.
             chunk_mode: bool = 'chunk',
+            scalar_output: bool = False,
     ) -> None:
         super().__init__()
 
@@ -77,6 +78,7 @@ def __init__(
             activation=activation,
             chunk_mode=chunk_mode) for _ in range(num_decoder_layers)])
         self.out_linear = nn.Linear(d_model, out_channels)
+        self.scalar_output = scalar_output
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         '''
@@ -89,10 +91,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         encoding = self.encoders(encoding)
 
         decoding = encoding
-        if self.pe is not None:  # position encoding
-            decoding = self.pe(decoding)
-        for layer in self.decoders:
-            decoding = layer(decoding, encoding)
+        if len(self.decoders):
+            if self.pe is not None:  # position encoding
+                decoding = self.pe(decoding)
+            for layer in self.decoders:
+                decoding = layer(decoding, encoding)
+
+        if self.scalar_output:  # if want scalar instead of seq output, take the first index from seq
+            decoding = decoding[:, 0, :]
 
         output = self.out_linear(decoding)
         return output

torcharc/module_builder.py

@@ -91,7 +91,7 @@ def infer_in_shape(arc: dict, xs: Union[torch.Tensor, NamedTuple]) -> None:
         arc.update(in_shape=in_shape)
     elif nn_type == 'FiLMMerge':
         assert ps.is_tuple(xs)
-        assert len(arc['in_names']) == 2, f'FiLMMerge in_names should only specify 2 keys for feature and conditioner'
+        assert len(arc['in_names']) == 2, 'FiLMMerge in_names should only specify 2 keys for feature and conditioner'
         shapes = {name: list(x.shape)[1:] for name, x in xs._asdict().items() if name in arc['in_names']}
         arc.update(shapes=shapes)
     else: