Merge pull request #48 from valentingol/conditional

🆙 Update to 2.0.2

Author: valentingol

.pylintrc

@@ -148,7 +148,8 @@ confidence=HIGH,
 disable=no-member,
         no-name-in-module,
         not-callable,
-        redefined-outer-name
+        redefined-outer-name,
+        duplicate-code,
 
 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option

configs/default/models/cond_sagan.yaml

@@ -6,11 +6,16 @@ data_size: 64  # could be either 32, 64, 128 or 256
 # NOTE: the number of blocks depends on the data size only:
 # data_size 32: 4 blocks, data_size 64: 5 blocks, data_size 128: 6 blocks, data_size 256: 7 blocks
 attn_layer_num: [3, 4]
-full_values: False
 z_dim: 128
 g_conv_dim: 64  # number of channels before the last layer of the generator
 d_conv_dim: 64  # number of channels after the first layer of the discriminator
 
 # cond_dim_ratio: define the number of channels of conditional feature maps
 # cond_dim = base_dim // cond_dim_ratio
 cond_dim_ratio: 8
+
+attention: !attention
+  n_heads: 1  # number of attention heads
+  out_layer: True  # whether to apply a linear layer to the output of the attention
+  qk_ratio: 8  # dimension for queries and keys is input_dim // qk_ratio
+  v_ratio: 2  # dimension for values is input_dim // qk_ratio, should be 1 if out_layer=False

configs/default/models/sagan.yaml

@@ -6,8 +6,13 @@ data_size: 64  # could be either 32, 64, 128 or 256
 # NOTE: the number of blocks depends on the data size only:
 # data_size 32: 4 blocks, data_size 64: 5 blocks, data_size 128: 6 blocks, data_size 256: 7 blocks
 attn_layer_num: [3, 4]
-full_values: False
 z_dim: 128
 g_conv_dim: 64  # number of channels before the last layer of the generator
 d_conv_dim: 64  # number of channels after the first layer of the discriminator
 cond_dim_ratio: -1 # no used here (only for conditional models)
+
+attention: !attention
+  n_heads: 1  # number of attention heads
+  out_layer: True  # whether to apply a linear layer to the output of the attention
+  qk_ratio: 8  # dimension for queries and keys is input_dim // qk_ratio
+  v_ratio: 2  # dimension for values is input_dim // qk_ratio, should be 1 if out_layer=False

configs/unittest/data32.yaml

@@ -15,8 +15,11 @@ model.attn_layer_num: [1, 2, 3]
 model.data_size: 32
 model.d_conv_dim: 8
 model.g_conv_dim: 8
-model.full_values: True
 model.init_method: orthogonal
+model.attention.n_heads: 1
+model.attention.out_layer: True
+model.attention.qk_ratio: 8
+model.attention.v_ratio: 2
 
 training.adv_loss: wgan-gp
 training.ema_start_step: 0

configs/unittest/data64.yaml

@@ -13,8 +13,11 @@ model.attn_layer_num: [4]
 model.data_size: 64
 model.d_conv_dim: 12
 model.g_conv_dim: 12
-model.full_values: False
 model.init_method: normal
+model.attention.n_heads: 4
+model.attention.out_layer: False
+model.attention.qk_ratio: 1
+model.attention.v_ratio: 1
 
 training.adv_loss: wgan-gp
 training.ema_start_step: 0

setup.py

@@ -5,7 +5,7 @@
 # Installation
 config = {
     'name': 'sagan-facies-modeling',
-    'version': '2.0.1',
+    'version': '2.0.2',
     'description': 'Facies modeling with SAGAN.',
     'author': 'Valentin Goldite',
     'author_email': '[email protected]',

tests/utils/conftest.py

@@ -4,9 +4,64 @@
 
 import numpy as np
 import pytest
+import torch
 from pytest_check import check_func
+from torch.utils.data import DataLoader
 
 from utils.configs import GlobalConfig
+from utils.data.data_loader import DistributedDataLoader
+
+
+class DataLoader64(DistributedDataLoader):
+    """Data loader for unit tests (data size 64)."""
+
+    def __init__(self) -> None:
+        # pylint: disable=super-init-not-called
+        self.n_classes = 4
+
+    def loader(self) -> DataLoader:
+        """Return pytorch data loader."""
+
+        class Dataset64(torch.utils.data.Dataset):
+            """Dataset for unit tests (data size 32)."""
+
+            def __getitem__(self, index: int) -> Tuple[torch.Tensor, int]:
+                return torch.randn(4, 64, 64), 0
+
+            def __len__(self) -> int:
+                return 10
+
+        return torch.utils.data.DataLoader(dataset=Dataset64(),
+                                           batch_size=2,
+                                           shuffle=True,
+                                           num_workers=0,
+                                           )
+
+
+class DataLoader32(DistributedDataLoader):
+    """Data loader for unit tests (data size 32)."""
+
+    def __init__(self) -> None:
+        # pylint: disable=super-init-not-called
+        self.n_classes = 4
+
+    def loader(self) -> DataLoader:
+        """Return pytorch data loader."""
+
+        class Dataset32(torch.utils.data.Dataset):
+            """Dataset for unit tests (data size 32)."""
+
+            def __getitem__(self, index: int) -> Tuple[torch.Tensor, int]:
+                return torch.randn(4, 32, 32), 0
+
+            def __len__(self) -> int:
+                return 10
+
+        return torch.utils.data.DataLoader(dataset=Dataset32(),
+                                           batch_size=2,
+                                           shuffle=True,
+                                           num_workers=0,
+                                           )
 
 
 @check_func

tests/utils/gan/cond_sagan/test_cond_modules.py

@@ -18,9 +18,9 @@ def test_sa_generator(configs: Tuple[GlobalConfig, GlobalConfig]) -> None:
     data, att_list = gen(z, pixel_maps, with_attn=True)
     check.equal(data.shape, (5, 4, 32, 32))
     check.equal(len(att_list), 3)
-    check.equal(att_list[0].shape, (5, 16, 16))
-    check.equal(att_list[1].shape, (5, 64, 64))
-    check.equal(att_list[2].shape, (5, 256, 256))
+    check.equal(att_list[0].shape, (5, 1, 16, 16))
+    check.equal(att_list[1].shape, (5, 1, 64, 64))
+    check.equal(att_list[2].shape, (5, 1, 256, 256))
     data = gen(z, pixel_maps, with_attn=False)
     check.is_instance(data, torch.Tensor)
 

tests/utils/gan/test_attention.py

@@ -3,16 +3,32 @@
 import pytest_check as check
 import torch
 
+from utils.configs import ConfigType
 from utils.gan.attention import SelfAttention
 
 
-def test_self_attention() -> None:
+def test_self_attention(configs: ConfigType) -> None:
     """Test SelfAttention."""
-    attention = SelfAttention(in_dim=16, att_dim=8, full_values=False)
+    config_32, config_64 = configs
+
+    # Case with out_layer = True
+    attention = SelfAttention(in_dim=16,
+                              attention_config=config_32.model.attention)
     out, attention = attention(torch.rand(2, 16, 9, 9))
     check.equal(out.shape, (2, 16, 9, 9))
-    check.equal(attention.shape, (2, 81, 81))
-    attention = SelfAttention(in_dim=16, att_dim=None, full_values=True)
+    check.equal(attention.shape, (2, 1, 81, 81))
+
+    # Case with out_layer = False
+    attention = SelfAttention(in_dim=16,
+                              attention_config=config_64.model.attention)
     out, attention = attention(torch.rand(2, 16, 9, 9))
     check.equal(out.shape, (2, 16, 9, 9))
-    check.equal(attention.shape, (2, 81, 81))
+    check.equal(attention.shape, (2, 4, 81, 81))
+
+    # Case with out_layer = False and v_ratio != 1
+    config_64_bis = config_64.copy()
+    config_64_bis.merge({"model": {"attention": {"v_ratio": 2}}},
+                        do_not_pre_process=True)
+    with check.raises(ValueError):
+        SelfAttention(in_dim=16,
+                      attention_config=config_64_bis.model.attention)

tests/utils/gan/test_base_trainer.py

@@ -3,7 +3,7 @@
 import os
 import os.path as osp
 import shutil
-from typing import Dict, List, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import ignite.distributed as idist
 import numpy as np
@@ -12,66 +12,12 @@
 from pytest_mock import MockerFixture
 from torch.nn import Module
 from torch.optim import Optimizer
-from torch.utils.data import DataLoader
 
-from tests.utils.conftest import check_exists
+from tests.utils.conftest import DataLoader32, DataLoader64, check_exists
 from utils.configs import GlobalConfig
-from utils.data.data_loader import DistributedDataLoader
 from utils.gan.base_trainer import BaseTrainerGAN, BatchType
 
 
-class DataLoader64(DistributedDataLoader):
-    """Data loader for unit tests (data size 64)."""
-
-    def __init__(self) -> None:
-        # pylint: disable=super-init-not-called
-        self.n_classes = 4
-
-    def loader(self) -> DataLoader:
-        """Return pytorch data loader."""
-
-        class Dataset64(torch.utils.data.Dataset):
-            """Dataset for unit tests (data size 32)."""
-
-            def __getitem__(self, index: int) -> Tuple[torch.Tensor, int]:
-                return torch.randn(4, 64, 64), 0
-
-            def __len__(self) -> int:
-                return 10
-
-        return torch.utils.data.DataLoader(dataset=Dataset64(),
-                                           batch_size=2,
-                                           shuffle=True,
-                                           num_workers=0,
-                                           )
-
-
-class DataLoader32(DistributedDataLoader):
-    """Data loader for unit tests (data size 32)."""
-
-    def __init__(self) -> None:
-        # pylint: disable=super-init-not-called
-        self.n_classes = 4
-
-    def loader(self) -> DataLoader:
-        """Return pytorch data loader."""
-
-        class Dataset32(torch.utils.data.Dataset):
-            """Dataset for unit tests (data size 32)."""
-
-            def __getitem__(self, index: int) -> Tuple[torch.Tensor, int]:
-                return torch.randn(4, 32, 32), 0
-
-            def __len__(self) -> int:
-                return 10
-
-        return torch.utils.data.DataLoader(dataset=Dataset32(),
-                                           batch_size=2,
-                                           shuffle=True,
-                                           num_workers=0,
-                                           )
-
-
 class TrainerTest(BaseTrainerGAN):
     """Test class for BaseTrainerGAN."""
     def train_generator(self, gen: Module, g_optimizer: Optimizer,
@@ -91,7 +37,7 @@ def train_discriminator(self, disc: Module, d_optimizer: Optimizer,
 
     def build_model_opt(self) -> Tuple[Module, Module, Optimizer, Optimizer]:
 
-        class Generator(torch.nn.Module):
+        class Generator(Module):
             """Simple Generator"""
             def __init__(self) -> None:
                 super().__init__()
@@ -148,7 +94,6 @@ def build_trainers() -> Tuple[BaseTrainerGAN, BaseTrainerGAN]:
 
 # Test TrainerSAGAN
 
-
 def test_init() -> None:
     """Test init method."""
     build_trainers()
@@ -174,8 +119,29 @@ def test_train(mocker: MockerFixture) -> None:
     np.save('tests/datasets/data32.npy', data32)
     np.save('tests/datasets/data64.npy', data64)
 
+    # Mock train_discriminator and train_generator
+    def side_effect_disc(disc: Module, d_optimizer: Optimizer,
+                         *args: Optional[Tuple], **kwargs: Optional[Dict]
+                         ) -> Tuple[Module, Optimizer,
+                                    Dict[str, torch.Tensor]]:
+        # pylint: disable=unused-argument
+        losses = {'d_loss': (torch.tensor(0.3), 'red', 6)}
+        return disc, d_optimizer, losses
+
+    def side_effect_gen(gen: Module, g_optimizer: Optimizer,
+                        *args: Optional[Tuple], **kwargs: Optional[Dict]
+                        ) -> Tuple[Module, Optimizer,
+                                   Dict[str, torch.Tensor]]:
+        # pylint: disable=unused-argument
+        losses = {'g_loss': (torch.tensor(0.2), 'green', 6)}
+        return gen, g_optimizer, losses
+
     trainers = build_trainers()
     for trainer in trainers:
+        trainer.train_discriminator = mocker.MagicMock(  # type: ignore
+            side_effect=side_effect_disc)
+        trainer.train_generator = mocker.MagicMock(  # type: ignore
+            side_effect=side_effect_gen)
         trainer.train()
         check_exists('res/tmp_test/models/generator_step_2.pth')
         check_exists('res/tmp_test/models/discriminator_step_2.pth')