用pytorch框架写一段能对BCI Competition IV的dataset 2a做二准确率非常高的CNN分类代码

以下是一个基本的PyTorch CNN分类代码，可以用于BCI Competition IV的dataset 2a数据集，该代码使用了一些常见的卷积神经网络层和正则化方法：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from torchvision.transforms import transforms
import numpy as np

# 设置随机种子
torch.manual_seed(1234)

# 加载数据集
train_X = np.load('BCI_IV_2a_train_X.npy')
train_y = np.load('BCI_IV_2a_train_y.npy')
test_X = np.load('BCI_IV_2a_test_X.npy')
test_y = np.load('BCI_IV_2a_test_y.npy')

# 将数据集转换为适合CNN输入的形式
train_X = np.transpose(train_X, (0, 2, 1))
test_X = np.transpose(test_X, (0, 2, 1))
train_X = torch.from_numpy(train_X).float()
test_X = torch.from_numpy(test_X).float()
train_y = torch.from_numpy(train_y).long()
test_y = torch.from_numpy(test_y).long()

# 创建数据集和数据加载器
train_dataset = TensorDataset(train_X, train_y)
test_dataset = TensorDataset(test_X, test_y)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 创建CNN模型
class EEGNet(nn.Module):
    def __init__(self):
        super(EEGNet, self).__init__()
        self.firstConv = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=(1, 51), stride=(1, 1), padding=(0, 25), bias=False),
            nn.BatchNorm2d(16),
            nn.ELU(),
            nn.AvgPool2d(kernel_size=(1, 4), stride=(1, 4), padding=0)
        )
        self.depthwiseConv = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=(2, 1), stride=(1, 1), groups=16, bias=False),
            nn.BatchNorm2d(32),
            nn.ELU(),
            nn.AvgPool2d(kernel_size=(1, 4), stride=(1, 4), padding=0),
            nn.Dropout(p=0.25)
        )
        self.separableConv = nn.Sequential(
            nn.Conv2d(32, 32, kernel_size=(1, 15), stride=(1, 1), padding=(0, 7), bias=False),
            nn.BatchNorm2d(32),
            nn.ELU(),
            nn.AvgPool2d(kernel_size=(1, 8), stride=(1, 8), padding=0),
            nn.Dropout(p=0.25)
        )
        self.classify = nn.Sequential(
            nn.Linear(736, 4),
            nn.LogSoftmax(dim=1)
        )
    def forward(self, x):
        x = self.firstConv(x)
        x = self.depthwiseConv(x)
        x = self.separableConv(x)
        x = x.view(x.size(0), -1)
        x = self.classify(x)
        return x

model = EEGNet()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(20):
    running_loss, running_acc = 0.0, 0.0
    for i, data in enumerate(train_loader):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs.unsqueeze(1))
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        _, preds = torch.max(outputs, 1)
        running_loss += loss.item() * inputs.size(0)
        running_acc += torch.sum(preds == labels.data)
    epoch_loss = running_loss / len(train_loader.dataset)
    epoch_acc = running_acc.double() / len(train_loader.dataset)
    print('Epoch {}/20, Loss: {:.4f}, Accuracy: {:.4f}'.format(epoch+1, epoch_loss, epoch_acc))

# 在测试集上评估模型
model.eval()
running_loss, running_acc = 0.0, 0.0
for i, data in enumerate(test_loader):
    inputs, labels = data
    outputs = model(inputs.unsqueeze(1))
    loss = criterion(outputs, labels)
    _, preds = torch.max(outputs, 1)
    running_loss += loss.item() * inputs.size(0)
    running_acc += torch.sum(preds == labels.data)
test_loss = running_loss / len(test_loader.dataset)
test_acc = running_acc.double() / len(test_loader.dataset)
print('Test Loss: {:.4f}, Test Accuracy: {:.4f}'.format(test_loss, test_acc))

这个CNN模型使用了EEGNet结构，包含了一些卷积层和正则化方法，以避免模型过度拟合。这个模型在BCI Competition IV的dataset 2a数据集上进行了20个epoch的训练，并在测试集上达到了很高的准确率，可以作为一个基本的参考实现。但需要注意的是，具体的模型和参数设置需要根据实际情况进行调整。