初始化zmq 项目

MI/Algorithm/conformer_2class.py

@@ -0,0 +1,409 @@
+"""
+EEG Conformer 
+
+Convolutional Transformer for EEG decoding
+
+Couple CNN and Transformer in a concise manner with amazing results
+"""
+# remember to change paths
+
+import os
+gpus = [0]
+os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
+os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, gpus))
+import numpy as np
+import math
+import random
+import time
+import datetime
+
+from torch.utils.data import DataLoader
+from torch.autograd import Variable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch import Tensor
+from einops import rearrange
+from einops.layers.torch import Rearrange, Reduce
+# from common_spatial_pattern import csp
+
+# from torch.utils.tensorboard import SummaryWriter
+from torch.backends import cudnn
+cudnn.benchmark = True
+cudnn.deterministic = True
+from sklearn.model_selection import train_test_split
+# writer = SummaryWriter('./TensorBoardX/')
+
+
+# Convolution module
+# use conv to capture local features, instead of postion embedding.
+class PatchEmbedding(nn.Module):
+    def __init__(self, emb_size=40,n_chan=8):
+        # self.patch_size = patch_size
+        super().__init__()
+
+        self.shallownet = nn.Sequential(
+            nn.Conv2d(1, 40, (1, 25), (1, 1)),
+            nn.Conv2d(40, 40, (n_chan, 1), (1, 1)),
+            nn.BatchNorm2d(40),
+            nn.ELU(),
+            nn.AvgPool2d((1, 75), (1, 15)),  # pooling acts as slicing to obtain 'patch' along the time dimension as in ViT
+            nn.Dropout(0.5),
+        )
+
+        self.projection = nn.Sequential(
+            nn.Conv2d(40, emb_size, (1, 1), stride=(1, 1)),  # transpose, conv could enhance fiting ability slightly
+            Rearrange('b e (h) (w) -> b (h w) e'),
+        )
+
+
+    def forward(self, x: Tensor) -> Tensor:
+        b, _, _, _ = x.shape
+        x = self.shallownet(x)
+        x = self.projection(x)
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, emb_size, num_heads, dropout):
+        super().__init__()
+        self.emb_size = emb_size
+        self.num_heads = num_heads
+        self.keys = nn.Linear(emb_size, emb_size)
+        self.queries = nn.Linear(emb_size, emb_size)
+        self.values = nn.Linear(emb_size, emb_size)
+        self.att_drop = nn.Dropout(dropout)
+        self.projection = nn.Linear(emb_size, emb_size)
+
+    def forward(self, x: Tensor, mask: Tensor = None) -> Tensor:
+        queries = rearrange(self.queries(x), "b n (h d) -> b h n d", h=self.num_heads)
+        keys = rearrange(self.keys(x), "b n (h d) -> b h n d", h=self.num_heads)
+        values = rearrange(self.values(x), "b n (h d) -> b h n d", h=self.num_heads)
+        energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys)  
+        if mask is not None:
+            fill_value = torch.finfo(torch.float32).min
+            energy.mask_fill(~mask, fill_value)
+
+        scaling = self.emb_size ** (1 / 2)
+        att = F.softmax(energy / scaling, dim=-1)
+        att = self.att_drop(att)
+        out = torch.einsum('bhal, bhlv -> bhav ', att, values)
+        out = rearrange(out, "b h n d -> b n (h d)")
+        out = self.projection(out)
+        return out
+
+
+class ResidualAdd(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        res = x
+        x = self.fn(x, **kwargs)
+        x += res
+        return x
+
+
+class FeedForwardBlock(nn.Sequential):
+    def __init__(self, emb_size, expansion, drop_p):
+        super().__init__(
+            nn.Linear(emb_size, expansion * emb_size),
+            nn.GELU(),
+            nn.Dropout(drop_p),
+            nn.Linear(expansion * emb_size, emb_size),
+        )
+
+
+class GELU(nn.Module):
+    def forward(self, input: Tensor) -> Tensor:
+        return input*0.5*(1.0+torch.erf(input/math.sqrt(2.0)))
+
+
+class TransformerEncoderBlock(nn.Sequential):
+    def __init__(self,
+                 emb_size,
+                 num_heads=10,
+                 drop_p=0.5,
+                 forward_expansion=4,
+                 forward_drop_p=0.5):
+        super().__init__(
+            ResidualAdd(nn.Sequential(
+                nn.LayerNorm(emb_size),
+                MultiHeadAttention(emb_size, num_heads, drop_p),
+                nn.Dropout(drop_p)
+            )),
+            ResidualAdd(nn.Sequential(
+                nn.LayerNorm(emb_size),
+                FeedForwardBlock(
+                    emb_size, expansion=forward_expansion, drop_p=forward_drop_p),
+                nn.Dropout(drop_p)
+            )
+            ))
+
+
+class TransformerEncoder(nn.Sequential):
+    def __init__(self, depth, emb_size):
+        super().__init__(*[TransformerEncoderBlock(emb_size) for _ in range(depth)])
+
+
+class ClassificationHead(nn.Sequential):
+    def __init__(self, emb_size, n_classes):
+        super().__init__()
+        
+        # global average pooling
+        self.clshead = nn.Sequential(
+            Reduce('b n e -> b e', reduction='mean'),
+            nn.LayerNorm(emb_size),
+            nn.Linear(emb_size, n_classes)
+        )
+        self.fc = nn.Sequential(
+            nn.Linear(2440, 256),
+            nn.ELU(),
+            nn.Dropout(0.5),
+            nn.Linear(256, 32),
+            nn.ELU(),
+            nn.Dropout(0.3),
+            nn.Linear(32, 2)
+        )
+
+    def forward(self, x):
+        x = x.contiguous().view(x.size(0), -1)
+        out = self.fc(x)
+        return out
+
+
+class Conformer(nn.Sequential):
+    def __init__(self, emb_size=40, depth=6, n_classes=2,n_chan=8, **kwargs):
+        super().__init__(
+
+            PatchEmbedding(emb_size,n_chan),
+            TransformerEncoder(depth, emb_size),
+            ClassificationHead(emb_size, n_classes)
+        )
+
+
+class ExP():
+    def __init__(self,n_chan):
+        super(ExP, self).__init__()
+        self.n_chan = n_chan
+        self.batch_size = 24
+        self.n_epochs = 250
+        self.c_dim = 4
+        self.lr = 0.0002
+        self.b1 = 0.5
+        self.b2 = 0.999
+
+        self.start_epoch = 0
+        # 创建目录
+        os.makedirs("online_Models", exist_ok=True)
+        self.log_write = open("./online_Models/log_result.txt", "w")
+
+
+        self.Tensor = torch.cuda.FloatTensor
+        self.LongTensor = torch.cuda.LongTensor
+
+        self.criterion_cls = torch.nn.CrossEntropyLoss().cuda()
+
+        self.model = Conformer(n_chan=self.n_chan).cuda()
+        self.model = nn.DataParallel(self.model, device_ids=[i for i in range(len(gpus))])
+        self.model = self.model.cuda()
+
+        # self.model = EEGNet().cuda()
+        # self.model = nn.DataParallel(self.model,device_ids=[i for i in range(len(gpus))])
+        # self.model = self.model.cuda()
+        # summary(self.model, (1, 8, 1000))
+
+
+    # Segmentation and Reconstruction (S&R) data augmentation
+    def interaug(self, timg, label):
+        # 确保输入是 numpy 数组（CPU）
+        if isinstance(timg, torch.Tensor):
+            timg = timg.cpu().numpy()
+        if isinstance(label, torch.Tensor):
+            label = label.cpu().numpy()
+
+        aug_data = []
+        aug_label = []
+        for cls4aug in range(2):
+            cls_idx = np.where(label == cls4aug + 1)
+            tmp_data = timg[cls_idx]
+            tmp_label = label[cls_idx]
+            tmp_aug_data = np.zeros((int(self.batch_size / 2), 1, self.n_chan, 1000))
+            for ri in range(int(self.batch_size / 2)):
+                for rj in range(8):
+                    rand_idx = np.random.randint(0, tmp_data.shape[0], 8)
+                    tmp_aug_data[ri, :, :, rj * 125:(rj + 1) * 125] = tmp_data[rand_idx[rj], :, :,
+                                                                      rj * 125:(rj + 1) * 125]
+
+            aug_data.append(tmp_aug_data)
+            aug_label.append(tmp_label[:int(self.batch_size / 2)])
+        aug_data = np.concatenate(aug_data)
+        aug_label = np.concatenate(aug_label)
+        aug_shuffle = np.random.permutation(len(aug_data))
+        aug_data = aug_data[aug_shuffle, :, :]
+        aug_label = aug_label[aug_shuffle]
+
+        # 返回 numpy 数组，由调用方决定是否移到 GPU
+        return aug_data, aug_label
+
+    def train(self,all_data,all_label,model_path):
+        all_data = np.array(all_data);all_label = np.array(all_label)
+        all_data = np.expand_dims(all_data, axis=1)
+        train_data, test_data, train_label, test_label = train_test_split(all_data, all_label, test_size=0.2,
+                                                                          random_state=42, stratify=all_label,shuffle=True)
+
+        # === 优化：一次性预生成增强数据，避免每个 batch 都重复计算 ===
+        aug_data, aug_label = self.interaug(train_data, train_label)
+        # 将原始数据和增强数据合并，再一起打乱
+        train_data_full = np.concatenate([train_data, aug_data], axis=0)
+        train_label_full = np.concatenate([train_label, aug_label], axis=0)
+        shuffle_idx = np.random.permutation(len(train_data_full))
+        train_data_full = train_data_full[shuffle_idx]
+        train_label_full = train_label_full[shuffle_idx]
+
+        img = torch.from_numpy(train_data_full)
+        label = torch.from_numpy(train_label_full-1)
+
+        dataset = torch.utils.data.TensorDataset(img, label)
+        self.dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=self.batch_size, shuffle=True)
+
+        test_data = torch.from_numpy(test_data)
+        test_label = torch.from_numpy(test_label-1)
+        test_dataset = torch.utils.data.TensorDataset(test_data, test_label)
+        self.test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=self.batch_size, shuffle=True)
+
+        # Optimizers
+        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr, betas=(self.b1, self.b2))
+
+        test_data = Variable(test_data.type(self.Tensor))
+        test_label = Variable(test_label.type(self.LongTensor))
+
+        bestAcc = 0
+        averAcc = 0
+        num = 0
+        Y_true = 0
+        Y_pred = 0
+
+        # Train the cnn model
+        for e in range(self.n_epochs):
+            # in_epoch = time.time()
+            self.model.train()
+            for i, (img, label) in enumerate(self.dataloader):
+
+                img = Variable(img.cuda().type(self.Tensor))
+                label = Variable(label.cuda().type(self.LongTensor))
+
+                outputs = self.model(img)
+
+                loss = self.criterion_cls(outputs, label) 
+
+                self.optimizer.zero_grad()
+                loss.backward()
+                self.optimizer.step()
+
+
+            # out_epoch = time.time()
+
+
+            # test process
+            if (e + 1) % 1 == 0:
+                self.model.eval()
+                Cls = self.model(test_data)
+
+                loss_test = self.criterion_cls(Cls, test_label)
+                y_pred = torch.max(Cls, 1)[1]
+                acc = float((y_pred == test_label).cpu().numpy().astype(int).sum()) / float(test_label.size(0))
+                train_pred = torch.max(outputs, 1)[1]
+                train_acc = float((train_pred == label).cpu().numpy().astype(int).sum()) / float(label.size(0))
+
+                print('Epoch:', e,
+                      '  Train loss: %.6f' % loss.detach().cpu().numpy(),
+                      '  Test loss: %.6f' % loss_test.detach().cpu().numpy(),
+                      '  Train accuracy %.6f' % train_acc,
+                      '  Test accuracy is %.6f' % acc)
+
+                self.log_write.write(str(e) + "    " + str(acc) + "\n")
+                num = num + 1
+                averAcc = averAcc + acc
+                if acc > bestAcc:
+                    bestAcc = acc
+                    Y_true = test_label
+                    Y_pred = y_pred
+
+
+        torch.save(self.model, model_path)
+        averAcc = averAcc / num
+        print('The average accuracy is:', averAcc)
+        print('The best accuracy is:', bestAcc)
+        self.log_write.write('The average accuracy is: ' + str(averAcc) + "\n")
+        self.log_write.write('The best accuracy is: ' + str(bestAcc) + "\n")
+
+        return bestAcc, averAcc, Y_true, Y_pred
+        # writer.close()
+
+
+def onlineTrain(data_queue,result_queue):
+    import torch
+    print(f"[DEBUG] torch.__version__ = {torch.__version__}")
+    print(f"[DEBUG] torch.cuda.is_available() = {torch.cuda.is_available()}")
+    if torch.cuda.is_available():
+        print(f"[DEBUG] GPU = {torch.cuda.get_device_name(0)}")
+    try:
+        starttime = datetime.datetime.now()
+
+        # seed_n = np.random.randint(2025)
+        seed_n = 1877
+        random.seed(seed_n)
+        np.random.seed(seed_n)
+        torch.manual_seed(seed_n)
+        torch.cuda.manual_seed(seed_n)
+        torch.cuda.manual_seed_all(seed_n)
+
+
+        # 从队列获取训练数据
+        data = data_queue.get(timeout=30)
+        all_data, all_label,model_path,n_chan = data['data'], data['label'],data['modelPath'],data['n_chan']
+        exp = ExP(n_chan)
+        print('训练参数： ',np.shape(all_data),np.shape(all_label),model_path)
+        bestAcc, averAcc, Y_true, Y_pred = exp.train(all_data,all_label,model_path)
+        print('THE BEST ACCURACY IS ' + str(bestAcc))
+
+        endtime = datetime.datetime.now()
+        print('train duration: ',str(endtime - starttime))
+
+        # 将模型或参数传回
+        result_queue.put({
+            'status': 'success',
+            'model_state': model_path,  # 或保存路径
+            'timestamp': time.time()
+        })
+    except Exception as e:
+        result_queue.put({'status': 'error', 'msg': str(e)})
+
+def offlineTrain(all_data,all_label,modelPath):
+    starttime = datetime.datetime.now()
+
+    # seed_n = np.random.randint(2025)
+    seed_n = 1877
+    print('seed is ' + str(seed_n))
+    random.seed(seed_n)
+    np.random.seed(seed_n)
+    torch.manual_seed(seed_n)
+    torch.cuda.manual_seed(seed_n)
+    torch.cuda.manual_seed_all(seed_n)
+
+    exp = ExP()
+
+    bestAcc, averAcc, Y_true, Y_pred = exp.train(all_data,all_label,modelPath)
+    print('THE BEST ACCURACY IS ' + str(bestAcc))
+
+    endtime = datetime.datetime.now()
+    print('train duration: ',str(endtime - starttime))
+
+
+
+if __name__ == "__main__":
+    print(time.asctime(time.localtime(time.time())))
+    print(time.asctime(time.localtime(time.time())))

MI/Algorithm/conformer_2class_cpu.py

@@ -0,0 +1,382 @@
+"""
+EEG Conformer 
+
+Convolutional Transformer for EEG decoding
+
+Couple CNN and Transformer in a concise manner with amazing results
+"""
+# remember to change paths
+import os
+import numpy as np
+import math
+import random
+import time
+import datetime
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torch import nn
+from torch import Tensor
+from einops import rearrange
+from einops.layers.torch import Rearrange, Reduce
+from torch.backends import cudnn
+from sklearn.model_selection import train_test_split
+# writer = SummaryWriter('./TensorBoardX/')
+
+
+# Convolution module
+# use conv to capture local features, instead of postion embedding.
+class PatchEmbedding(nn.Module):
+    def __init__(self, emb_size=40):
+        # self.patch_size = patch_size
+        super().__init__()
+
+        self.shallownet = nn.Sequential(
+            nn.Conv2d(1, 40, (1, 25), (1, 1)),
+            nn.Conv2d(40, 40, (8, 1), (1, 1)),
+            nn.BatchNorm2d(40),
+            nn.ELU(),
+            nn.AvgPool2d((1, 75), (1, 15)),  # pooling acts as slicing to obtain 'patch' along the time dimension as in ViT
+            nn.Dropout(0.5),
+        )
+
+        self.projection = nn.Sequential(
+            nn.Conv2d(40, emb_size, (1, 1), stride=(1, 1)),  # transpose, conv could enhance fiting ability slightly
+            Rearrange('b e (h) (w) -> b (h w) e'),
+        )
+
+
+    def forward(self, x: Tensor) -> Tensor:
+        b, _, _, _ = x.shape
+        x = self.shallownet(x)
+        x = self.projection(x)
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, emb_size, num_heads, dropout):
+        super().__init__()
+        self.emb_size = emb_size
+        self.num_heads = num_heads
+        self.keys = nn.Linear(emb_size, emb_size)
+        self.queries = nn.Linear(emb_size, emb_size)
+        self.values = nn.Linear(emb_size, emb_size)
+        self.att_drop = nn.Dropout(dropout)
+        self.projection = nn.Linear(emb_size, emb_size)
+
+    def forward(self, x: Tensor, mask: Tensor = None) -> Tensor:
+        queries = rearrange(self.queries(x), "b n (h d) -> b h n d", h=self.num_heads)
+        keys = rearrange(self.keys(x), "b n (h d) -> b h n d", h=self.num_heads)
+        values = rearrange(self.values(x), "b n (h d) -> b h n d", h=self.num_heads)
+        energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys)  
+        if mask is not None:
+            fill_value = torch.finfo(torch.float32).min
+            energy.mask_fill(~mask, fill_value)
+
+        scaling = self.emb_size ** (1 / 2)
+        att = F.softmax(energy / scaling, dim=-1)
+        att = self.att_drop(att)
+        out = torch.einsum('bhal, bhlv -> bhav ', att, values)
+        out = rearrange(out, "b h n d -> b n (h d)")
+        out = self.projection(out)
+        return out
+
+
+class ResidualAdd(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        res = x
+        x = self.fn(x, **kwargs)
+        x += res
+        return x
+
+
+class FeedForwardBlock(nn.Sequential):
+    def __init__(self, emb_size, expansion, drop_p):
+        super().__init__(
+            nn.Linear(emb_size, expansion * emb_size),
+            nn.GELU(),
+            nn.Dropout(drop_p),
+            nn.Linear(expansion * emb_size, emb_size),
+        )
+
+
+class GELU(nn.Module):
+    def forward(self, input: Tensor) -> Tensor:
+        return input*0.5*(1.0+torch.erf(input/math.sqrt(2.0)))
+
+
+class TransformerEncoderBlock(nn.Sequential):
+    def __init__(self,
+                 emb_size,
+                 num_heads=10,
+                 drop_p=0.5,
+                 forward_expansion=4,
+                 forward_drop_p=0.5):
+        super().__init__(
+            ResidualAdd(nn.Sequential(
+                nn.LayerNorm(emb_size),
+                MultiHeadAttention(emb_size, num_heads, drop_p),
+                nn.Dropout(drop_p)
+            )),
+            ResidualAdd(nn.Sequential(
+                nn.LayerNorm(emb_size),
+                FeedForwardBlock(
+                    emb_size, expansion=forward_expansion, drop_p=forward_drop_p),
+                nn.Dropout(drop_p)
+            )
+            ))
+
+
+class TransformerEncoder(nn.Sequential):
+    def __init__(self, depth, emb_size):
+        super().__init__(*[TransformerEncoderBlock(emb_size) for _ in range(depth)])
+
+
+class ClassificationHead(nn.Sequential):
+    def __init__(self, emb_size, n_classes):
+        super().__init__()
+        
+        # global average pooling
+        self.clshead = nn.Sequential(
+            Reduce('b n e -> b e', reduction='mean'),
+            nn.LayerNorm(emb_size),
+            nn.Linear(emb_size, n_classes)
+        )
+        self.fc = nn.Sequential(
+            nn.Linear(2440, 256),
+            nn.ELU(),
+            nn.Dropout(0.5),
+            nn.Linear(256, 32),
+            nn.ELU(),
+            nn.Dropout(0.3),
+            nn.Linear(32, 2)
+        )
+
+    def forward(self, x):
+        x = x.contiguous().view(x.size(0), -1)
+        out = self.fc(x)
+        return out
+
+
+class Conformer(nn.Sequential):
+    def __init__(self, emb_size=40, depth=6, n_classes=2, **kwargs):
+        super().__init__(
+
+            PatchEmbedding(emb_size),
+            TransformerEncoder(depth, emb_size),
+            ClassificationHead(emb_size, n_classes)
+        )
+
+
+class ExP():
+    def __init__(self):
+        super(ExP, self).__init__()
+        self.batch_size = 24
+        self.n_epochs = 250
+        self.c_dim = 4
+        self.lr = 0.0002
+        self.b1 = 0.5
+        self.b2 = 0.999
+
+        self.start_epoch = 0
+
+        self.log_write = open("./online_Models/log_result.txt", "w")
+
+        # 自动选择设备：有 GPU 用 GPU，否则用 CPU
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        # self.device = torch.device("cpu")
+        print(f"Using device: {self.device}")
+
+        # 定义张量类型（不再强制使用 cuda）
+        self.Tensor = torch.FloatTensor
+        self.LongTensor = torch.LongTensor
+
+        # 将模型移到指定设备
+        self.model = Conformer().to(self.device)
+
+        # 损失函数也移到设备
+        self.criterion_cls = torch.nn.CrossEntropyLoss().to(self.device)
+
+        # self.model = EEGNet().cuda()
+        # self.model = nn.DataParallel(self.model,device_ids=[i for i in range(len(gpus))])
+        # self.model = self.model.cuda()
+        # summary(self.model, (1, 8, 1000))
+
+
+    # Segmentation and Reconstruction (S&R) data augmentation
+    def interaug(self, timg, label):
+        aug_data = []
+        aug_label = []
+        for cls4aug in range(2):
+            cls_idx = np.where(label == cls4aug + 1)
+            tmp_data = timg[cls_idx]
+            tmp_label = label[cls_idx]
+            tmp_aug_data = np.zeros((int(self.batch_size / 2), 1, 8, 1000))
+            for ri in range(int(self.batch_size / 2)):
+                for rj in range(8):
+                    rand_idx = np.random.randint(0, tmp_data.shape[0], 8)
+                    tmp_aug_data[ri, :, :, rj * 125:(rj + 1) * 125] = tmp_data[rand_idx[rj], :, :,
+                                                                      rj * 125:(rj + 1) * 125]
+
+            aug_data.append(tmp_aug_data)
+            aug_label.append(tmp_label[:int(self.batch_size / 2)])
+        aug_data = np.concatenate(aug_data)
+        aug_label = np.concatenate(aug_label)
+        aug_shuffle = np.random.permutation(len(aug_data))
+        aug_data = aug_data[aug_shuffle, :, :]
+        aug_label = aug_label[aug_shuffle]
+
+        aug_data = torch.from_numpy(aug_data).float().to(self.device)
+        aug_label = torch.from_numpy(aug_label - 1).long().to(self.device)
+        return aug_data, aug_label
+
+    def train(self,all_data,all_label,model_path):
+        all_data = np.array(all_data);all_label = np.array(all_label)
+        all_data = np.expand_dims(all_data, axis=1)
+        train_data, test_data, train_label, test_label = train_test_split(all_data, all_label, test_size=0.2,
+                                                                          random_state=42, stratify=all_label,shuffle=True)
+        # 转为 Tensor
+        img = torch.from_numpy(train_data).float().to(self.device)
+        label = torch.from_numpy(train_label - 1).long().to(self.device)
+
+        dataset = torch.utils.data.TensorDataset(img, label)
+        self.dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=self.batch_size, shuffle=True)
+
+        test_data = torch.from_numpy(test_data).float().to(self.device)
+        test_label = torch.from_numpy(test_label - 1).long().to(self.device)
+        test_dataset = torch.utils.data.TensorDataset(test_data, test_label)
+        self.test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=self.batch_size, shuffle=True)
+
+        # Optimizers
+        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr, betas=(self.b1, self.b2))
+
+        bestAcc = 0
+        averAcc = 0
+        num = 0
+        Y_true = 0
+        Y_pred = 0
+
+        # Train the cnn model
+        for e in range(self.n_epochs):
+            # in_epoch = time.time()
+            self.model.train()
+            for i, (img, label) in enumerate(self.dataloader):
+
+                # data augmentation
+                aug_data, aug_label = self.interaug(train_data, train_label)
+                img = torch.cat((img, aug_data))
+                label = torch.cat((label, aug_label))
+
+
+                outputs = self.model(img)
+
+                loss = self.criterion_cls(outputs, label) 
+
+                self.optimizer.zero_grad()
+                loss.backward()
+                self.optimizer.step()
+
+
+            # out_epoch = time.time()
+
+
+            # test process
+            if (e + 1) % 1 == 0:
+                self.model.eval()
+                with torch.no_grad():
+                    Cls = self.model(test_data)
+
+                    loss_test = self.criterion_cls(Cls, test_label)
+                    y_pred = torch.max(Cls, 1)[1]
+                    acc = float((y_pred == test_label).cpu().numpy().astype(int).sum()) / float(test_label.size(0))
+                    train_pred = torch.max(outputs, 1)[1]
+                    train_acc = float((train_pred == label).cpu().numpy().astype(int).sum()) / float(label.size(0))
+
+                print('Epoch:', e,
+                      '  Train loss: %.6f' % loss.detach().cpu().numpy(),
+                      '  Test loss: %.6f' % loss_test.detach().cpu().numpy(),
+                      '  Train accuracy %.6f' % train_acc,
+                      '  Test accuracy is %.6f' % acc)
+
+                self.log_write.write(str(e) + "    " + str(acc) + "\n")
+                num = num + 1
+                averAcc = averAcc + acc
+                if acc > bestAcc:
+                    bestAcc = acc
+                    Y_true = test_label
+                    Y_pred = y_pred
+
+
+        torch.save(self.model, model_path)
+        averAcc = averAcc / num
+        print('The average accuracy is:', averAcc)
+        print('The best accuracy is:', bestAcc)
+        self.log_write.write('The average accuracy is: ' + str(averAcc) + "\n")
+        self.log_write.write('The best accuracy is: ' + str(bestAcc) + "\n")
+
+        return bestAcc, averAcc, Y_true, Y_pred
+        # writer.close()
+
+
+def onlineTrain(data_queue,result_queue):
+    try:
+        starttime = datetime.datetime.now()
+
+        # seed_n = np.random.randint(2025)
+        seed_n = 1877
+        random.seed(seed_n)
+        np.random.seed(seed_n)
+        torch.manual_seed(seed_n)
+        torch.cuda.manual_seed(seed_n)
+        torch.cuda.manual_seed_all(seed_n)
+
+        exp = ExP()
+        # 从队列获取训练数据
+        data = data_queue.get(timeout=30)
+        all_data, all_label,model_path = data['data'], data['label'],data['modelPath']
+        print('训练参数： ',np.shape(all_data),np.shape(all_label),model_path)
+        bestAcc, averAcc, Y_true, Y_pred = exp.train(all_data,all_label,model_path)
+        print('THE BEST ACCURACY IS ' + str(bestAcc))
+
+        endtime = datetime.datetime.now()
+        print('train duration: ',str(endtime - starttime))
+
+        # 将模型或参数传回
+        result_queue.put({
+            'status': 'success',
+            'model_state': model_path,  # 或保存路径
+            'timestamp': time.time()
+        })
+    except Exception as e:
+        result_queue.put({'status': 'error', 'msg': str(e)})
+
+def offlineTrain(all_data,all_label,modelPath):
+    starttime = datetime.datetime.now()
+
+    # seed_n = np.random.randint(2025)
+    seed_n = 1877
+    print('seed is ' + str(seed_n))
+    random.seed(seed_n)
+    np.random.seed(seed_n)
+    torch.manual_seed(seed_n)
+    torch.cuda.manual_seed(seed_n)
+    torch.cuda.manual_seed_all(seed_n)
+
+    exp = ExP()
+
+    bestAcc, averAcc, Y_true, Y_pred = exp.train(all_data,all_label,modelPath)
+    print('THE BEST ACCURACY IS ' + str(bestAcc))
+
+    endtime = datetime.datetime.now()
+    print('train duration: ',str(endtime - starttime))
+
+
+
+if __name__ == "__main__":
+    print(time.asctime(time.localtime(time.time())))
+    print(time.asctime(time.localtime(time.time())))

MI/Algorithm/otherModels.py

@@ -0,0 +1,184 @@
+from torchsummary import summary
+import torch
+import torch.nn as nn
+
+
+def weights_init(m):
+    if isinstance(m, nn.Conv2d):
+        nn.init.xavier_uniform_(m.weight)
+        # nn.init.constant(m.bias, 0)  # bias may be none
+
+    elif isinstance(m, nn.BatchNorm2d):
+        nn.init.constant_(m.weight, 1)
+        nn.init.constant_(m.bias, 0)
+
+    elif isinstance(m, nn.Linear):
+        nn.init.xavier_uniform_(m.weight)
+        nn.init.constant_(m.bias, 0)
+
+
+
+def square_activation(x):
+    return torch.square(x)
+
+
+def safe_log(x):
+    return torch.clip(torch.log(x), min=1e-7, max=1e7)
+
+
+class ShallowConvNet(nn.Module):
+    def __init__(self, num_classes=3, chans=19, samples=768):
+        super(ShallowConvNet, self).__init__()
+        self.conv_nums = 40
+        self.features = nn.Sequential(
+            nn.Conv2d(1, self.conv_nums, (1, 25)),
+            nn.Conv2d(self.conv_nums, self.conv_nums, (chans, 1), bias=False),
+            nn.BatchNorm2d(self.conv_nums)
+        )
+        self.avgpool = nn.AvgPool2d(kernel_size=(1, 75), stride=(1, 15))
+        self.dropout = nn.Dropout()
+
+        out = torch.ones((1, 1, chans, samples))
+        out = self.features(out)
+        out = self.avgpool(out)
+        n_out_time = out.cpu().data.numpy().shape
+        self.classifier = nn.Linear(n_out_time[-1] * n_out_time[-2] * n_out_time[-3], num_classes)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = square_activation(x)
+        x = self.avgpool(x)
+        x = safe_log(x)
+        x = self.dropout(x)
+
+        features = torch.flatten(x, 1)
+        cls = self.classifier(features)
+        return cls
+
+
+class EEGNet(nn.Module):
+    def __init__(self, num_classes=2, chans=8, samples=1000, dropout_rate=0.5, kernel_length=64, F1=8,
+                 F2=16,):
+        super(EEGNet, self).__init__()
+
+        self.features = nn.Sequential(
+            nn.Conv2d(1, F1, kernel_size=(1, kernel_length), bias=False),
+            nn.BatchNorm2d(F1),
+            nn.Conv2d(F1, F1, kernel_size=(chans, 1), groups=F1, bias=False),  # groups=F1 for depthWiseConv
+            nn.BatchNorm2d(F1),
+            nn.ELU(inplace=True),
+            # nn.ReLU(),
+            nn.AvgPool2d((1, 4)),
+            nn.Dropout(dropout_rate),
+            # for SeparableCon2D
+            # SeparableConv2D(F1, F2, kernel1_size=(1, 16), bias=False),
+            nn.Conv2d(F1, F1, kernel_size=(1, 16), groups=F1, bias=False),  # groups=F1 for depthWiseConv
+            nn.BatchNorm2d(F1),
+            nn.ELU(inplace=True),
+            # nn.ReLU(),
+            nn.Conv2d(F1, F2, kernel_size=(1, 1), groups=1, bias=False),  # point-wise cnn
+            nn.BatchNorm2d(F2),
+            # nn.ReLU(),
+            nn.ELU(inplace=True),
+            nn.AvgPool2d((1, 8)),
+            nn.Dropout(p=dropout_rate),
+            # nn.Dropout(p=0.5),
+        )
+        out = torch.ones((1, 1, chans, samples))
+        out = self.features(out)
+        n_out_time = out.cpu().data.numpy().shape
+        self.classifier = nn.Linear(n_out_time[-1] * n_out_time[-2] * n_out_time[-3], num_classes)
+
+    def forward(self, x):
+        conv_features = self.features(x)
+        features = torch.flatten(conv_features, 1)
+        cls = self.classifier(features)
+        return cls
+
+
+class LMDA(nn.Module):
+    """
+    LMDA-Net for the paper
+    """
+    def __init__(self, chans=19, samples=768, num_classes=3, depth=9, kernel=75, channel_depth1=24, channel_depth2=9,
+                ave_depth=1, avepool=5):
+        super(LMDA, self).__init__()
+        self.ave_depth = ave_depth
+        self.channel_weight = nn.Parameter(torch.randn(depth, 1, chans), requires_grad=True)
+        nn.init.xavier_uniform_(self.channel_weight.data)
+
+
+        self.time_conv = nn.Sequential(
+            nn.Conv2d(depth, channel_depth1, kernel_size=(1, 1), groups=1, bias=False),
+            nn.BatchNorm2d(channel_depth1),
+            nn.Conv2d(channel_depth1, channel_depth1, kernel_size=(1, kernel),
+                      groups=channel_depth1, bias=False),
+            nn.BatchNorm2d(channel_depth1),
+            nn.GELU(),
+        )
+        # self.avgPool1 = nn.AvgPool2d((1, 24))
+        self.chanel_conv = nn.Sequential(
+            nn.Conv2d(channel_depth1, channel_depth2, kernel_size=(1, 1), groups=1, bias=False),
+            nn.BatchNorm2d(channel_depth2),
+            nn.Conv2d(channel_depth2, channel_depth2, kernel_size=(chans, 1), groups=channel_depth2, bias=False),
+            nn.BatchNorm2d(channel_depth2),
+            nn.GELU(),
+        )
+
+        self.norm = nn.Sequential(
+            nn.AvgPool3d(kernel_size=(1, 1, avepool)),
+            # nn.AdaptiveAvgPool3d((9, 1, 35)),
+            nn.Dropout(p=0.65),
+        )
+
+        # 定义自动填充模块
+        out = torch.ones((1, 1, chans, samples))
+        out = torch.einsum('bdcw, hdc->bhcw', out, self.channel_weight)
+        out = self.time_conv(out)
+        out = self.chanel_conv(out)
+        out = self.norm(out)
+        n_out_time = out.cpu().data.numpy().shape
+        print('In ShallowNet, n_out_time shape: ', n_out_time)
+        self.classifier = nn.Linear(n_out_time[-1]*n_out_time[-2]*n_out_time[-3], num_classes)
+
+    def EEGDepthAttention(self, x):
+        # x: input features with shape [N, C, H, W]
+
+        N, C, H, W = x.size()
+        # K = W if W % 2 else W + 1
+        k = 7
+        adaptive_pool = nn.AdaptiveAvgPool2d((1, W))
+        conv = nn.Conv2d(1, 1, kernel_size=(k, 1), padding=(k//2, 0), bias=True).to(x.device)  # original kernel k
+        nn.init.xavier_uniform_(conv.weight)
+        nn.init.constant_(conv.bias, 0)
+        softmax = nn.Softmax(dim=-2)
+        x_pool = adaptive_pool(x)
+        x_transpose = x_pool.transpose(-2, -3)
+        y = conv(x_transpose)
+        y = softmax(y)
+        y = y.transpose(-2, -3)
+        return y * C * x
+
+    def forward(self, x):
+        x = torch.einsum('bdcw, hdc->bhcw', x, self.channel_weight)
+
+        x_time = self.time_conv(x)  # batch, depth1, channel, samples_
+        x_time = self.EEGDepthAttention(x_time)  # DA1
+
+        x = self.chanel_conv(x_time)  # batch, depth2, 1, samples_
+        x = self.norm(x)
+
+        features = torch.flatten(x, 1)
+        cls = self.classifier(features)
+        return cls
+
+
+if __name__ == '__main__':
+    model = ShallowConvNet(num_classes=4, chans=22, samples=1125).cuda()
+    a = torch.randn(12, 1, 3, 875).cuda().float()
+    l2 = model(a)
+    model_optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
+    summary(model, show_input=True)
+
+    print(l2.shape)
+