27、ResNet50处理STEW数据集，用于情感三分类+全备的代码

1、数据介绍

IEEE-Datasets-STEW:SIMULTANEOUS TASK EEG WORKLOAD DATASET ：

该数据集由48名受试者的原始EEG数据组成，他们参加了利用SIMKAP多任务测试进行的多任务工作负荷实验。受试者在休息时的大脑活动也在测试前被记录下来，也包括在其中。Emotiv EPOC设备，采样频率为128Hz，有14个通道，用于获取数据，每个案例都有2.5分钟的EEG记录。受试者还被要求在每个阶段后以1到9的评分标准对其感知的心理工作量进行评分，评分结果在单独的文件中提供。

说明：每个受试者的数据遵循命名惯例：subno_task.txt。例如，sub01_lo.txt将是受试者1在休息时的原始脑电数据，而sub23_hi.txt将是受试者23在多任务测试中的原始脑电数据。每个数据文件的行对应于记录中的样本，列对应于EEG设备的14个通道： AF3, F7, F3, FC5, T7, P7, O1, O2, P8, T8, FC6, F4, F8, AF4。

数据说明、下载地址：

STEW: Simultaneous Task EEG Workload Data Set | IEEE Journals & Magazine | IEEE Xplore

2、代码

本次使用ResNet50，去做此情感数据的分类工作，数据导入+模型训练+测试代码如下：

import torch
import torchvision.datasets
from torch.utils.data import Dataset        # 继承Dataset类
import os
from PIL import Image
import numpy as np
from torchvision import transforms
 
 
# 预处理
data_transform = transforms.Compose([
    transforms.Resize((224,224)),           # 缩放图像
    transforms.ToTensor(),                  # 转为Tenso
    transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))       # 标准化
])


path =  r'C:\STEW\test'

for root,dirs,files in os.walk(path):
        print('root',root) #遍历到该目录地址
        print('dirs',dirs) #遍历到该目录下的子目录名 []
        print('files',files)  #遍历到该目录下的文件  []
 
def read_txt_files(path):
    # 创建文件名列表
    file_names = []
    # 遍历给定目录及其子目录下的所有文件
    for root, dirs, files in os.walk(path):
        # 遍历所有文件
        for file in files:
            # 如果是 .txt 文件，则加入文件名列表
            if file.endswith('.txt'): # endswith () 方法用于判断字符串是否以指定后缀结尾，如果以指定后缀结尾返回True，否则返回False。
                file_names.append(os.path.join(root, file))
    # 返回文件名列表
    return file_names

class DogCat(Dataset):      # 数据处理
    def __init__(self,root,transforms = None):                  # 初始化，指定路径,是否预处理等等
 
        #['cat.15454.jpg', 'cat.445.jpg', 'cat.46456.jpg', 'cat.656165.jpg', 'dog.123.jpg', 'dog.15564.jpg', 'dog.4545.jpg', 'dog.456465.jpg']
        imgs = os.listdir(root)
 
        self.imgs = [os.path.join(root,img) for img in imgs]    # 取出root下所有的文件
        self.transforms = data_transform                        # 图像预处理
 
    def __getitem__(self, index):       # 读取图片
        img_path = self.imgs[index]
        label = 1 if 'dog' in img_path.split('/')[-1] else 0 
        #然后，就可以根据每个路径的id去做label了。将img_path 路径按照 '/ '分割，-1代表取最后一个字符串，如果里面有dog就为1，cat就为0.
 
        data = Image.open(img_path)
 
        if self.transforms:     # 图像预处理
            data = self.transforms(data)
 
        return data,label
 
    def __len__(self):
        return len(self.imgs)
 
dataset = DogCat('./data/',transforms=True)
 
for img,label in dataset:
    print('img:',img.size(),'label:',label)
'''
img: torch.Size([3, 224, 224]) label: 0
img: torch.Size([3, 224, 224]) label: 0
img: torch.Size([3, 224, 224]) label: 0
img: torch.Size([3, 224, 224]) label: 0
img: torch.Size([3, 224, 224]) label: 1
img: torch.Size([3, 224, 224]) label: 1
img: torch.Size([3, 224, 224]) label: 1
img: torch.Size([3, 224, 224]) label: 1
'''

import os
 
# 获取file_path路径下的所有TXT文本内容和文件名
def get_text_list(file_path):
    files = os.listdir(file_path)
    text_list = []
    for file in files:
        with open(os.path.join(file_path, file), "r", encoding="UTF-8") as f:
            text_list.append(f.read())
    return text_list, files
 
class ImageFolderCustom(Dataset):

    # 2. Initialize with a targ_dir and transform (optional) parameter
    def __init__(self, targ_dir: str, transform=None) -> None:

        # 3. Create class attributes
        # Get all image paths
        self.paths = list(pathlib.Path(targ_dir).glob("*/*.jpg")) # note: you'd have to update this if you've got .png's or .jpeg's
        # Setup transforms
        self.transform = transform
        # Create classes and class_to_idx attributes
        self.classes, self.class_to_idx = find_classes(targ_dir)

    # 4. Make function to load images
    def load_image(self, index: int) -> Image.Image:
        "Opens an image via a path and returns it."
        image_path = self.paths[index]
        return Image.open(image_path) 

    # 5. Overwrite the __len__() method (optional but recommended for subclasses of torch.utils.data.Dataset)
    def __len__(self) -> int:
        "Returns the total number of samples."
        return len(self.paths)

    # 6. Overwrite the __getitem__() method (required for subclasses of torch.utils.data.Dataset)
    def __getitem__(self, index: int) -> Tuple[torch.Tensor, int]:
        "Returns one sample of data, data and label (X, y)."
        img = self.load_image(index)
        class_name  = self.paths[index].parent.name # expects path in data_folder/class_name/image.jpeg
        class_idx = self.class_to_idx[class_name]

        # Transform if necessary
        if self.transform:
            return self.transform(img), class_idx # return data, label (X, y)
        else:
            return img, class_idx # return data, label (X, y)
                  
import torchvision as tv
import numpy as np
import torch
import time
import os
from torch import nn, optim
from torchvision.models import resnet50
from torchvision.transforms import transforms
 
os.environ["CUDA_VISIBLE_DEVICE"] = "0,1,2"
 
# cifar-10进行测验

class Cutout(object):
    """Randomly mask out one or more patches from an image.
    Args:
        n_holes (int): Number of patches to cut out of each image.
        length (int): The length (in pixels) of each square patch.
    """
    def __init__(self, n_holes, length):
        self.n_holes = n_holes
        self.length = length
 
    def __call__(self, img):
        """
        Args:
            img (Tensor): Tensor image of size (C, H, W).
        Returns:
            Tensor: Image with n_holes of dimension length x length cut out of it.
        """
        h = img.size(1)
        w = img.size(2)
 
        mask = np.ones((h, w), np.float32)
 
        for n in range(self.n_holes):
            y = np.random.randint(h)
            x = np.random.randint(w)
 
            y1 = np.clip(y - self.length // 2, 0, h)
            y2 = np.clip(y + self.length // 2, 0, h)
            x1 = np.clip(x - self.length // 2, 0, w)
            x2 = np.clip(x + self.length // 2, 0, w)
 
            mask[y1: y2, x1: x2] = 0.
 
        mask = torch.from_numpy(mask)
        mask = mask.expand_as(img)
        img = img * mask
 
        return img
 
def load_data_cifar10(batch_size=128,num_workers=2):
    # 操作合集
    # Data augmentation
    train_transform_1 = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
        transforms.RandomRotation(degrees=(-80,80)),  # 随机角度翻转
        transforms.ToTensor(),
        transforms.Normalize(
            (0.491339968,0.48215827,0.44653124), (0.24703233,0.24348505,0.26158768)  # 两者分别为(mean,std)
        ),
        Cutout(1, 16),  # 务必放在ToTensor的后面
    ])
    train_transform_2 = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(
            (0.491339968, 0.48215827, 0.44653124), (0.24703233, 0.24348505, 0.26158768)  # 两者分别为(mean,std)
        )
    ])
    test_transform = transforms.Compose([
        transforms.Resize((224,224)),
        transforms.ToTensor(),
        transforms.Normalize(
            (0.491339968, 0.48215827, 0.44653124), (0.24703233, 0.24348505, 0.26158768)  # 两者分别为(mean,std)
        )
    ])
    # 训练集1
    trainset1 = tv.datasets.CIFAR10(
        root='data',
        train=True,
        download=False,
        transform=train_transform_1,
    )
    # 训练集2
    trainset2 = tv.datasets.CIFAR10(
        root='data',
        train=True,
        download=False,
        transform=train_transform_2,
    )
    # 测试集
    testset = tv.datasets.CIFAR10(
        root='data',
        train=False,
        download=False,
        transform=test_transform,
    )
    # 训练数据加载器1
    trainloader1 = torch.utils.data.DataLoader(
        trainset1,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=(torch.cuda.is_available())
    )
    # 训练数据加载器2
    trainloader2 = torch.utils.data.DataLoader(
        trainset2,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=(torch.cuda.is_available())
    )
    # 测试数据加载器
    testloader = torch.utils.data.DataLoader(
        testset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=(torch.cuda.is_available())
    )
 
    return trainloader1,trainloader2,testloader
 
def main():
    start = time.time()
    batch_size = 128
    cifar_train1,cifar_train2,cifar_test = load_data_cifar10(batch_size=batch_size)
    model = resnet50().cuda()
    # model.load_state_dict(torch.load('_ResNet50.pth'))
    # 存在已保存的参数文件
    # model = nn.DataParallel(model,device_ids=[0,])  # 又套一层
    model = nn.DataParallel(model,device_ids=[0,1,2])
    loss = nn.CrossEntropyLoss().cuda()
    optimizer = optim.Adam(model.parameters(),lr=0.001)
    for epoch in range(50):
        model.train()  # 训练时务必写
        loss_=0.0
        num=0.0
        # train on trainloader1(data augmentation) and trainloader2
        for i,data in enumerate(cifar_train1,0):
            x, label = data
            x, label = x.cuda(),label.cuda()
            # x
            p = model(x) #output
            l = loss(p,label) #loss
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            loss_ += float(l.mean().item())
            num+=1
        for i, data in enumerate(cifar_train2, 0):
            x, label = data
            x, label = x.cuda(), label.cuda()
            # x
            p = model(x)
            l = loss(p, label)
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            loss_ += float(l.mean().item())
            num += 1
        model.eval()  # 评估时务必写
        print("loss:",float(loss_)/num)
        # test on trainloader2,testloader
        with torch.no_grad():
            total_correct = 0
            total_num = 0
            for x, label in cifar_train2:
                # [b, 3, 32, 32]
                # [b]
                x, label = x.cuda(), label.cuda()
                # [b, 10]
                logits = model(x)
                # [b]
                pred = logits.argmax(dim=1)
                # [b] vs [b] => scalar tensor
                correct = torch.eq(pred, label).float().sum().item()
                total_correct += correct
                total_num += x.size(0)
                # print(correct)
            acc_1 = total_correct / total_num
        # Test
        with torch.no_grad():
            total_correct = 0
            total_num = 0
            for x, label in cifar_test:
                # [b, 3, 32, 32]
                # [b]
                x, label = x.cuda(), label.cuda()
                # [b, 10]
                logits = model(x) #output
                # [b]
                pred = logits.argmax(dim=1)
                # [b] vs [b] => scalar tensor
                correct = torch.eq(pred, label).float().sum().item()
                total_correct += correct
                total_num += x.size(0)
                # print(correct)
            acc_2 = total_correct / total_num
            print(epoch+1,'train acc',acc_1,'|','test acc:', acc_2)
    # 保存时只保存model.module
    torch.save(model.module.state_dict(),'resnet50.pth')
    print("The interval is :",time.time() - start)
 
 
if __name__ == '__main__':
    main()