pytorch训练自己的数据集（完整版）

标签制作部分

我们要想用深度学习去训练我们的神经网络模型，首先，我们要告诉我们的电脑，我们的每一张图片，分别代表什么，对于这种问题，我们通过制作一个txt文档去解决，好，那我们现在就去构建一下我们txt文件，具体代码如下：

import os
import copy

def write_txt(path,txt_path):
    num =len(os.listdir(path))
    file_path = txt_path
    file = open (file_path, 'w')
    Y = 0
    c = os.listdir(path)
    for category in c:
        C = c.index(category)
        for imgs in os.listdir(os.path.join(path,category)):
            file.write(category+'/'+imgs+'|'+ str(Y)+ '\n')
        Y = Y+1
    return

if __name__ == "__main__":
    write_txt('C:/Users/Elegantmadman/Desktop/data/cats_and_dogs_filtered/validation/',
              'C:/Users/Elegantmadman/Desktop/data/cats_and_dogs_filtered/7.txt')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

上述代码是非常简单的，我们只需要提供根目录，就能够获得我们根目录下的分类文件夹，进而对我们的分类文件夹进行自动编号，然后获得文件夹下的图片名，我们将两者之间用“|”符号进行标记，以方便我们在后续进行按标识分割。

训练

我们在进行训练我们的数据集时，有两点问题需要解决，一是解决数据的读取问题，一是解决模型搭建问题。下面我们先来解决一下数据读取问题，这需要我们定义自己的Data类。代码如下：

class MyData(torch.utils.data.Dataset):
    def __init__(self, root, datatxt, transform=None, target_transform=None):
        super(MyData, self).__init__()
        file_txt = open(datatxt,'r')
        imgs = []
        for line in file_txt:
            line = line.rstrip()
            words = line.split('|')
            imgs.append((words[0], words[1]))

        self.imgs = imgs
        self.root = root
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):
        random.shuffle(self.imgs)
        name, label = self.imgs[index]
        img = Image.open(self.root + name).convert('RGB')
        if self.transform is not None:
            img = self.transform(img)
        label = int(label)
        #label_tensor = torch.Tensor([0,0])
        #label_tensor[label]=1
        return img, label

    def __len__(self):
        return len(self.imgs)

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

再去定义我们的模型

lass Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)
        self.pool = nn.MaxPool2d((2, 2))
        self.pool1 = nn.MaxPool2d((2, 2))
        self.conv2 = nn.Conv2d(16, 32, 3)
        self.fc1 = nn.Linear(36*36*32, 120)
        self.fc2 = nn.Linear(120, 60)
        self.fc3 = nn.Linear(60, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool1(F.relu(self.conv2(x)))
        x = x.view(-1, 36*36*32)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

下面进行完整的训练代码展示：

import torch
import random
from PIL import Image
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

class MyData(torch.utils.data.Dataset):
    def __init__(self, root, datatxt, transform=None, target_transform=None):
        super(MyData, self).__init__()
        file_txt = open(datatxt,'r')
        imgs = []
        for line in file_txt:
            line = line.rstrip()
            words = line.split('|')
            imgs.append((words[0], words[1]))

        self.imgs = imgs
        self.root = root
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):
        random.shuffle(self.imgs)
        name, label = self.imgs[index]
        img = Image.open(self.root + name).convert('RGB')
        if self.transform is not None:
            img = self.transform(img)
        label = int(label)
        #label_tensor = torch.Tensor([0,0])
        #label_tensor[label]=1
        return img, label

    def __len__(self):
        return len(self.imgs)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)
        self.pool = nn.MaxPool2d((2, 2))
        self.pool1 = nn.MaxPool2d((2, 2))
        self.conv2 = nn.Conv2d(16, 32, 3)
        self.fc1 = nn.Linear(36*36*32, 120)
        self.fc2 = nn.Linear(120, 60)
        self.fc3 = nn.Linear(60, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool1(F.relu(self.conv2(x)))
        x = x.view(-1, 36*36*32)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

classes = ['猫','狗']
transform = transforms.Compose(
    [transforms.Resize((152, 152)),transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_data=MyData(root ='C:/Users/Elegantmadman/Desktop/data/cats_and_dogs_filtered/train/',
                  datatxt='C:/Users/Elegantmadman/Desktop/data/cats_and_dogs_filtered/'+'6.txt',
                  transform=transform)
test_data=MyData(root ='C:/Users/Elegantmadman/Desktop/data/cats_and_dogs_filtered/validation/',
                 datatxt='C:/Users/Elegantmadman/Desktop/data/cats_and_dogs_filtered/'+'7.txt',
                 transform=transform)
train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=4, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_data, batch_size=4)

def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

dataiter = iter(train_loader)
images, labels = dataiter.next()

imshow(torchvision.utils.make_grid(images))
print(labels)

net = Net()

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(2):
    running_loss = 0.0
    for i, data in enumerate(train_loader, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = net(inputs)
        #print(outputs)
        #print(labels)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 200 == 0:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 200))
            running_loss = 0.0

print('Finished Training')

PATH = 'C:/Users/Elegantmadman/Desktop//cifar_net.pth'
torch.save(net.state_dict(), PATH)

dataiter = iter(test_loader)
images, labels = dataiter.next()

# print images
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))

outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]
                              for j in range(4)))


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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

测试与应用

下面我们将我们的模型应用和测试

net = Net()
net.load_state_dict(torch.load(PATH))

outputs = net(images)
correct = 0
total = 0
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

好的，至此我们就完成了所有的工作。