scitex 2.0.0__py2.py3-none-any.whl

scitex/nn/_BNet.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2023-05-15 16:44:27 (ywatanabe)"
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchsummary import summary
+import scitex
+import numpy as np
+import scitex
+
+
+class BHead(nn.Module):
+    def __init__(self, n_chs_in, n_chs_out):
+        super().__init__()
+        self.sa = scitex.nn.SpatialAttention(n_chs_in)
+        self.conv11 = nn.Conv1d(
+            in_channels=n_chs_in, out_channels=n_chs_out, kernel_size=1
+        )
+
+    def forward(self, x):
+        x = self.sa(x)
+        x = self.conv11(x)
+        return x
+
+
+class BNet(nn.Module):
+    def __init__(self, BNet_config, MNet_config):
+        super().__init__()
+        self.dummy_param = nn.Parameter(torch.empty(0))
+        N_VIRTUAL_CHS = 32
+
+        self.sc = scitex.nn.SwapChannels()
+        self.dc = scitex.nn.DropoutChannels(dropout=0.01)
+        self.fgc = scitex.nn.FreqGainChanger(
+            BNet_config["n_bands"], BNet_config["SAMP_RATE"]
+        )
+        self.heads = nn.ModuleList(
+            [
+                BHead(n_ch, N_VIRTUAL_CHS).to(self.dummy_param.device)
+                for n_ch in BNet_config["n_chs"]
+            ]
+        )
+
+        self.cgcs = [scitex.nn.ChannelGainChanger(n_ch) for n_ch in BNet_config["n_chs"]]
+        # self.cgc = scitex.nn.ChannelGainChanger(N_VIRTUAL_CHS)
+
+        MNet_config["n_chs"] = N_VIRTUAL_CHS  # BNet_config["n_chs"] # override
+        self.MNet = scitex.nn.MNet_1000(MNet_config)
+
+        self.fcs = nn.ModuleList(
+            [
+                nn.Sequential(
+                    # nn.Linear(N_FC_IN, config["n_fc1"]),
+                    nn.Mish(),
+                    nn.Dropout(BNet_config["d_ratio1"]),
+                    nn.Linear(BNet_config["n_fc1"], BNet_config["n_fc2"]),
+                    nn.Mish(),
+                    nn.Dropout(BNet_config["d_ratio2"]),
+                    nn.Linear(BNet_config["n_fc2"], BNet_config["n_classes"][i_head]),
+                )
+                for i_head, _ in enumerate(range(len(BNet_config["n_chs"])))
+            ]
+        )
+
+    @staticmethod
+    def _znorm_along_the_last_dim(x):
+        return (x - x.mean(dim=-1, keepdims=True)) / x.std(dim=-1, keepdims=True)
+
+    def forward(self, x, i_head):
+        x = self._znorm_along_the_last_dim(x)
+        # x = self.sc(x)
+        x = self.dc(x)
+        x = self.fgc(x)
+        x = self.cgcs[i_head](x)
+        x = self.heads[i_head](x)
+        import ipdb
+
+        ipdb.set_trace()
+        # x = self.cgc(x)
+        x = self.MNet.forward_bb(x)
+        x = self.fcs[i_head](x)
+        return x
+
+
+# BNet_config = {
+#     "n_chs": 32,
+#     "n_bands": 6,
+#     "SAMP_RATE": 1000,
+# }
+BNet_config = {
+    "n_bands": 6,
+    "SAMP_RATE": 250,
+    # "n_chs": 270,
+    "n_fc1": 1024,
+    "d_ratio1": 0.85,
+    "n_fc2": 256,
+    "d_ratio2": 0.85,
+}
+
+
+if __name__ == "__main__":
+    ## Demo data
+    # MEG
+    BS, N_CHS, SEQ_LEN = 16, 160, 1000
+    x_MEG = torch.rand(BS, N_CHS, SEQ_LEN).cuda()
+    # EEG
+    BS, N_CHS, SEQ_LEN = 16, 19, 1000
+    x_EEG = torch.rand(BS, N_CHS, SEQ_LEN).cuda()
+
+    # model = MNetBackBorn(scitex.nn.MNet_config).cuda()
+    # model(x_MEG)
+    # Model
+    BNet_config["n_chs"] = [160, 19]
+    BNet_config["n_classes"] = [2, 4]
+    model = BNet(BNet_config, scitex.nn.MNet_config).cuda()
+
+    # MEG
+    y = model(x_MEG, 0)
+    y = model(x_EEG, 1)
+
+    # # EEG
+    # y = model(x_EEG)
+
+    y.sum().backward()

scitex/nn/_BNet_Res.py

@@ -0,0 +1,164 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2023-05-15 17:09:58 (ywatanabe)"
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchsummary import summary
+import scitex
+import numpy as np
+import scitex
+
+
+class BHead(nn.Module):
+    def __init__(self, n_chs_in, n_chs_out):
+        super().__init__()
+        self.sa = scitex.nn.SpatialAttention(n_chs_in)
+        self.conv11 = nn.Conv1d(
+            in_channels=n_chs_in, out_channels=n_chs_out, kernel_size=1
+        )
+
+    def forward(self, x):
+        x = self.sa(x)
+        x = self.conv11(x)
+        return x
+
+
+class BNet(nn.Module):
+    def __init__(self, BNet_config, MNet_config):
+        super().__init__()
+        self.dummy_param = nn.Parameter(torch.empty(0))
+        # N_VIRTUAL_CHS = 32
+        # "n_virtual_chs":16,
+
+        self.sc = scitex.nn.SwapChannels()
+        self.dc = scitex.nn.DropoutChannels(dropout=0.01)
+        self.fgc = scitex.nn.FreqGainChanger(
+            BNet_config["n_bands"], BNet_config["SAMP_RATE"]
+        )
+        self.heads = nn.ModuleList(
+            [
+                BHead(n_ch, BNet_config["n_virtual_chs"]).to(self.dummy_param.device)
+                for n_ch in BNet_config["n_chs_of_modalities"]
+            ]
+        )
+
+        self.cgcs = [
+            scitex.nn.ChannelGainChanger(n_ch)
+            for n_ch in BNet_config["n_chs_of_modalities"]
+        ]
+        # self.cgc = scitex.nn.ChannelGainChanger(N_VIRTUAL_CHS)
+
+        # MNet_config["n_chs"] = BNet_config["n_virtual_chs"]  # BNet_config["n_chs"] # override
+
+        n_chs = BNet_config["n_virtual_chs"]
+        self.blk1 = scitex.nn.ResNetBasicBlock(n_chs, n_chs)
+        self.blk2 = scitex.nn.ResNetBasicBlock(int(n_chs / 2**1), int(n_chs / 2**1))
+        self.blk3 = scitex.nn.ResNetBasicBlock(int(n_chs / 2**2), int(n_chs / 2**2))
+        self.blk4 = scitex.nn.ResNetBasicBlock(int(n_chs / 2**3), int(n_chs / 2**3))
+        self.blk5 = scitex.nn.ResNetBasicBlock(1, 1)
+        self.blk6 = scitex.nn.ResNetBasicBlock(1, 1)
+        self.blk7 = scitex.nn.ResNetBasicBlock(1, 1)
+
+        # self.MNet = scitex.nn.MNet_1000(MNet_config)
+
+        # self.fcs = nn.ModuleList(
+        #     [
+        #         nn.Sequential(
+        #             # nn.Linear(N_FC_IN, config["n_fc1"]),
+        #             nn.Mish(),
+        #             nn.Dropout(BNet_config["d_ratio1"]),
+        #             nn.Linear(BNet_config["n_fc1"], BNet_config["n_fc2"]),
+        #             nn.Mish(),
+        #             nn.Dropout(BNet_config["d_ratio2"]),
+        #             nn.Linear(BNet_config["n_fc2"], BNet_config["n_classes_of_modalities"][i_head]),
+        #         )
+        #         for i_head, _ in enumerate(range(len(BNet_config["n_chs_of_modalities"])))
+        #     ]
+        # )
+
+    @staticmethod
+    def _znorm_along_the_last_dim(x):
+        return (x - x.mean(dim=-1, keepdims=True)) / x.std(dim=-1, keepdims=True)
+
+    def forward(self, x, i_head):
+        x = self._znorm_along_the_last_dim(x)
+        # x = self.sc(x)
+        x = self.dc(x)
+        x = self.fgc(x)
+        x = self.cgcs[i_head](x)
+        x = self.heads[i_head](x)
+
+        x = self.blk1(x)
+        x = F.avg_pool1d(x.transpose(1, 2), kernel_size=2).transpose(1, 2)
+        x = F.avg_pool1d(x, kernel_size=2)
+        x = self.blk2(x)
+        x = F.avg_pool1d(x.transpose(1, 2), kernel_size=2).transpose(1, 2)
+        x = F.avg_pool1d(x, kernel_size=2)
+        x = self.blk3(x)
+        x = F.avg_pool1d(x.transpose(1, 2), kernel_size=2).transpose(1, 2)
+        x = F.avg_pool1d(x, kernel_size=2)
+        x = self.blk4(x)
+        x = F.avg_pool1d(x.transpose(1, 2), kernel_size=2).transpose(1, 2)
+        x = F.avg_pool1d(x, kernel_size=2)
+
+        x = self.blk5(x)
+        x = F.avg_pool1d(x, kernel_size=2)
+        x = self.blk6(x)
+        x = F.avg_pool1d(x, kernel_size=2)
+        x = self.blk7(x)
+        x = F.avg_pool1d(x, kernel_size=2)
+
+        import ipdb
+
+        ipdb.set_trace()
+
+        # x = self.cgc(x)
+        x = self.MNet.forward_bb(x)
+        x = self.fcs[i_head](x)
+        return x
+
+
+# BNet_config = {
+#     "n_chs": 32,
+#     "n_bands": 6,
+#     "SAMP_RATE": 1000,
+# }
+BNet_config = {
+    "n_bands": 6,
+    "n_virtual_chs": 16,
+    "SAMP_RATE": 250,
+    "n_fc1": 1024,
+    "d_ratio1": 0.85,
+    "n_fc2": 256,
+    "d_ratio2": 0.85,
+}
+
+
+if __name__ == "__main__":
+    ## Demo data
+    # MEG
+    BS, N_CHS, SEQ_LEN = 16, 160, 1024
+    x_MEG = torch.rand(BS, N_CHS, SEQ_LEN).cuda()
+    # EEG
+    BS, N_CHS, SEQ_LEN = 16, 19, 1024
+    x_EEG = torch.rand(BS, N_CHS, SEQ_LEN).cuda()
+
+    # m = scitex.nn.ResNetBasicBlock(19, 19).cuda()
+    # m(x_EEG)
+    # model = MNetBackBorn(scitex.nn.MNet_config).cuda()
+    # model(x_MEG)
+    # Model
+    BNet_config["n_chs_of_modalities"] = [160, 19]
+    BNet_config["n_classes_of_modalities"] = [2, 4]
+    model = BNet(BNet_config, scitex.nn.MNet_config).cuda()
+
+    # MEG
+    y = model(x_MEG, 0)
+    y = model(x_EEG, 1)
+
+    # # EEG
+    # y = model(x_EEG)
+
+    y.sum().backward()

scitex/nn/_ChannelGainChanger.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2023-04-23 11:02:45 (ywatanabe)"
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchsummary import summary
+import scitex
+import numpy as np
+
+
+class ChannelGainChanger(nn.Module):
+    def __init__(
+        self,
+        n_chs,
+    ):
+        super().__init__()
+        self.n_chs = n_chs
+
+    def forward(self, x):
+        """x: [batch_size, n_chs, seq_len]"""
+        if self.training:
+            ch_gains = (
+                torch.rand(self.n_chs).unsqueeze(0).unsqueeze(-1).to(x.device) + 0.5
+            )
+            ch_gains = F.softmax(ch_gains, dim=1)
+            x *= ch_gains
+
+        return x
+
+
+if __name__ == "__main__":
+    ## Demo data
+    bs, n_chs, seq_len = 16, 360, 1000
+    x = torch.rand(bs, n_chs, seq_len)
+
+    cgc = ChGainChanger(n_chs)
+    print(cgc(x).shape)  # [16, 19, 1000]
+
+    # sb = SubjectBlock(n_chs=n_chs)
+    # print(sb(x, s).shape) # [16, 270, 1000]
+
+    # summary(sb, x, s)

scitex/nn/_DropoutChannels.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2023-05-04 21:50:22 (ywatanabe)"
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchsummary import summary
+import scitex
+import numpy as np
+import random
+
+
+class DropoutChannels(nn.Module):
+    def __init__(self, dropout=0.5):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x):
+        """x: [batch_size, n_chs, seq_len]"""
+        if self.training:
+            orig_chs = torch.arange(x.shape[1])
+
+            indi_orig = self.dropout(torch.ones(x.shape[1])).bool()
+            chs_to_shuffle = orig_chs[~indi_orig]
+
+            x[:, chs_to_shuffle] = torch.randn(x[:, chs_to_shuffle].shape).to(x.device)
+
+            # rand_chs = random.sample(list(np.array(chs_to_shuffle)), len(chs_to_shuffle))
+
+            # swapped_chs = orig_chs.clone()
+            # swapped_chs[~indi_orig] = torch.LongTensor(rand_chs)
+
+            # x = x[:, swapped_chs.long(), :]
+
+        return x
+
+
+if __name__ == "__main__":
+    ## Demo data
+    bs, n_chs, seq_len = 16, 360, 1000
+    x = torch.rand(bs, n_chs, seq_len)
+
+    dc = DropoutChannels(dropout=0.1)
+    print(dc(x).shape)  # [16, 19, 1000]
+
+    # sb = SubjectBlock(n_chs=n_chs)
+    # print(sb(x, s).shape) # [16, 270, 1000]
+
+    # summary(sb, x, s)