machine_learning_projects/3.宝可梦数据集分析/torch实现CNN服饰分类.py

import os
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

# 配置训练环境和超参数
# 根据系统的可用设备选择将张量放到GPU1或CPU上进行运算
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')

# 配置其他超参数
batch_size = 256  # 每个批次训练的样本数
num_workers = 0
learning_rate = 1e-4
epochs = 20  # 迭代次数

# 数据的读入和加载
from torchvision import transforms

image_size = 28
data_transform = transforms.Compose([
    transforms.ToPILImage(),  # 将张量转化为PIL图像对象
    transforms.Resize(image_size),
    transforms.ToTensor()  # 将PIL图像转化为张量
])

# 读入csv格式的数据，自行构建Dataset类
class FMDataset(Dataset):
    def __init__(self, df, transform=None):  # 默认没有数据转换操作
        self.df = df
        self.transform = transform
        self.images = df.iloc[:, 1:].values.astype(np.uint8)  # .values将提取到的DataFrame数据转化为(无符号八位整数)numpy数组
        self.labels = df.iloc[:, 0].values

    def __len__(self):  # 魔术方法，用于定义类的行为和操作，以模拟内置类型或实现类的特定功能
        return len(self.images)  # 这里用于返回图片数据集的长度

    def __getitem__(self, idx):  # 用于通过索引来访问数据集和标签
        image = self.images[idx].reshape(28, 28, 1)
        label = int(self.labels[idx])  # 根据索引获取对应对象
        if self.transform is not None:
            image = self.transform(image)
        else:
            image = torch.tensor(image/255., dtype=torch.float)  # 归一化
        label = torch.tensor(label, dtype=torch.long)
        return image, label  # 返回图像和标签元组


train_df = pd.read_csv('fashion-mnist_train.csv')
test_df = pd.read_csv('fashion-mnist_test.csv')
train_data = FMDataset(train_df, data_transform)
test_data = FMDataset(test_df, data_transform)

# 定义DataLoader类，以便在训练和测试时加载数据。
# DataLoader类是pytorch自带的类，将数据集封装为可迭代的数据加载器
train_loader = DataLoader(train_data, batch_size=batch_size,
                          shuffle=True, num_workers=num_workers,  # shuffle=True：在每个epoch开始时对数据集进行随机重排
                          drop_last=True)  # 如果最后一个批次样本数不足将被丢弃
test_loader = DataLoader(test_data, batch_size=batch_size,
                         shuffle=False, num_workers=num_workers)

# 可视化操作，用于验证读入的数据是否正确
import matplotlib.pyplot as plt

image, label = next(iter(train_loader))  # iter()将train_loader转化为一个迭代器对象，next()用于获得下一个批次的数据
print(image.shape, label.shape)
plt.imshow(image[0][0], cmap='gray')  # matplotlib中用于显示图像的函数
plt.show()


# 手搭CNN网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2),
            nn.Dropout(0.3),
            nn.Conv2d(32, 64, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, stride=2),
            nn.Dropout(0.3)
        )
        self.fc = nn.Sequential(
            nn.Linear(64 * 4 * 4, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def forward(self, x):
        x = self.conv(x)
        x = x.view(-1, 64 * 4 * 4)
        x = self.fc(x)
        return x


model = Net()
model = model.cpu()

# 设定损失函数
# torch.nn模块自带交叉熵损失
criterion = nn.CrossEntropyLoss()

# 设定优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练和测试
def train(epoch):
    model.train()
    train_loss = 0
    for data, label in train_loader:
        data, label = data.cpu(), label.cpu()
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, label)
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * data.size(0)
    train_loss = train_loss / len(train_loader.dataset)
    print('Epoch: {}\tTraining Loss: {:.4f}'.format(epoch, train_loss))

def val(epoch):
    model.eval()
    val_loss = 0
    gt_labels = []
    pred_labels = []
    with torch.no_grad():
        for data, label in test_loader:
            data, label = data.cpu(), label.cpu()
            output = model(data)
            preds = torch.argmax(output, 1)
            gt_labels.append(preds.cpu().data.numpy())
            pred_labels.append(preds.cpu().data.numpy())
            loss = criterion(output, label)
            val_loss += loss.item() * data.size(0)
    val_loss = val_loss / len(test_loader.dataset)
    gt_labels, pred_labels = np.concatenate(gt_labels), np.concatenate(pred_labels)
    acc = np.sum(gt_labels == pred_labels) / len(pred_labels)
    print('Epoch: {} \tValidation Loss: {:.4f}, Accuracy:{:.4f}'.format(epoch, val_loss, acc))


for epoch in range(1, epochs + 1):
    train(epoch)
    val(epoch)