pytorch模型实例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

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


# 配置 GPU
# 方案一
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# 方案二，后续对要使用 GPU 的变量用 .to(device) 即可
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

# 配置超参数
batch_size = 256
num_workers = 4    # 对于 Windows 用户，这里应设置为 0，否则会出现【多线程错误】
lr = 1e-4
epochs = 20

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

# 数据读入和加载
# 自己构建 Dataset，同时还需要对数据进行必要的变换
# 比如说需要将图片统一为一致的大小，需要将数据格式转为Tensor类，等等
# 可以借助 torchvision 包来完成
from torchvision import transforms, datasets


# 数据变换
image_size = 28
data_transform = transforms.Compose([
    transforms.ToPILImage(),    # 如果使用内置数据集读取方式则不需要
    transforms.Resize(image_size),
    transforms.ToTensor()
])

# 数据读取方式一：使用 torchvision 自带数据集，下载可能需要一段时间
train_data = datasets.FashionMNIST(
    root='./',
    train=True,
    download=True,
    transform=data_transform
)
test_data = datasets.FashionMNIST(
    root='./',
    train=False,
    download=True,
    transform=data_transform
)

# 读取方式二：读入 csv 格式的数据，自行构建 Dataset 类
# csv 数据下载链接：https://www.kaggle.com/zalando-research/fashionmnist
class FMDataset(Dataset):
    def __init__(self, df, transform=None):
        self.df = df
        self.transform = transform
        # 注意下面的 .astype(np.uint8)
        self.images = df.iloc[:, 1:].values.astype(np.uint8)
        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("./FashionMNIST/fashion-mnist_train.csv")
test_df = pd.read_csv("./FashionMNIST/fashion-mnist_test.csv")
train_data = FMDataset(train_df, data_transform)
test_data = FMDataset(test_df, data_transform)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

# 定义 DataLoader 类，用于在训练和测试时加载数据
train_loader = DataLoader(
    train_data,
    batch_size=batch_size,
    shuffle=True,
    num_workers=num_workers,
    drop_last=True
)
test_loader = DataLoader(
    test_data,
    batch_size=batch_size,
    shuffle=False,
    num_workers=num_workers
)

1
2
3
4
5
6
7

# 数据可视化操作
import matplotlib.pyplot as plt


image, label = next(iter(train_loader))
print(image.shape, label.shape)
plt.imshow(image[0][0], cmap="gray")

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

# 模型搭建
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)
        # x = nn.functional.normalize(x)
        return x


model = Net()
model = model.cuda()

1
2
3
4
5
6

# 损失函数
# PyTorch 会自动把整数型的 label 转为 one-hot 型，用于计算交叉熵损失
# 但是需要确保 label 是从 0 开始的
# 同时模型不加 softmax 层（因为使用了 logits 计算）

criterion = nn.CrossEntropyLoss()

1
2
3

# 优化器

optimizer = optim.Adam(model.parameters(), lr=0.001)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

# 训练和测试（验证）
# 各自封装成函数，方便后续调用
def train(epoch):
    model.train()
    train_loss = 0
    for data, label in train_loader:
        data, label = data.cuda(), label.cuda()
        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: {:.6f}'.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.cuda(), label.cuda()
            output = model(data)
            preds = torch.argmax(output, 1)
            gt_labels.append(label.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: {:.6f}, Accuracy: {:6f}'.format(epoch, val_loss, acc))

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

1
2
3
4

# 模型保存

save_path = "./FahionModel.pkl"
torch.save(model, save_path)