[人工智能]pytorch单物体图像分类

学习时不应该只关注算法模型的搭建，更应该关注怎么读取数据。

代码是使用datasets.ImageFolder来读取图片数据的，后续要自己学习怎么将图片数据写入txt,然后通过txt文件来获取我们需要的数据

1.代码知识对模型进行训练，并在测试集上检验模型的准确率。并没有用训练好的模型来预测一张图片的代码，后续需要添加。

2.在最开始的写代码的时候，会报一个错
“OSError: broken data stream when reading image file”
不知道是什么原因，添加了下面的代码之后就没问题了。

from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

下面是完整的代码： 我的数据集是放在hymenoptera_data这个文件夹下面的，完全可以自己找图片放进文件夹中，文件夹的格式如下：
（代码中只用到了train和valid这两个文件夹，然后每个文件夹里面就是物体种类的目录，本代码只有两类，所以train和valid里面只有两个目录）

import torch
from torch.autograd import Variable
import torchvision
from torchvision import datasets, transforms, models
import os
import matplotlib.pyplot as plt
import time
import torch.optim as optim

from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

# 1.数据加载
data_dir = './hymenoptera_data'
# 1.1 定义要对数据进行的处理
data_transform = {x: transforms.Compose([transforms.Resize([224, 224]), transforms.ToTensor()]) for x in ["train", "valid"]}
image_datasets = {x: datasets.ImageFolder(root=os.path.join(data_dir, x),  transform=data_transform[x]) for x in ["train", "valid"]}
dataloader = {x: torch.utils.data.DataLoader(dataset=image_datasets[x], batch_size=7, shuffle=True) for x in ["train", "valid"]}
train_data, val_data = dataloader["train"], dataloader["valid"]

index_classes = image_datasets["train"].class_to_idx
print(index_classes)
example_classes = image_datasets["train"].classes
print(example_classes)


# 2.数据预览, 在训练的时候可以注释掉
# X_example, y_example = next(iter(dataloader["train"]))
# img = torchvision.utils.make_grid(X_example)
# img = img.numpy().transpose([1, 2, 0])
# for i in range(len(y_example)):
#     index = y_example[i]
#     print(example_classes[index], end='   ')
#     if (i+1)%8 == 0:
#         print()
# plt.imshow(img)
# plt.show()


# 3.模型加载, 并对模型进行微调
net = models.resnet18(pretrained=True)
fc_features = net.fc.in_features
num_classes = 2
net.fc = torch.nn.Linear(fc_features, num_classes)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 4.pytorch fine tune 微调(冻结一部分层)
for i, param in enumerate(net.parameters()):
    if i < 30:
        param.requires_grad = False
net.to(device)

# 5.定义损失函数，以及优化器
LR = 1e-3
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=LR)


def train(epoch, loss_list):
    running_loss = 0
    for batch_idx, (data, target) in enumerate(train_data):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        out = net(data)
        loss = criterion(out, target)
        loss.backward()
        optimizer.step()

        loss_list.append(loss.item())
        running_loss += loss.item()
        if batch_idx % 10 == 9:
            print('[%d, %5d] loss:%.3f' % (epoch + 1, batch_idx + 1, running_loss / 3))
            running_loss = 0
    return loss_list


def test():
    correct, total = 0, 0
    with torch.no_grad():
        for _, (data, target) in enumerate(val_data):
            data, target = data.to(device), target.to(device)
            out = net(data)
            prediction =torch.max(out.data, dim=1)[1]
            total += target.size(0)
            correct += (prediction==target).sum().item()
        print('Accuracy on test set: (%d/%d)=%d %%' % (correct, total, 100 * correct / total))



if __name__ == '__main__':
    loss_list = []
    for epoch in range(20):
        train(epoch, loss_list)
        test()

    torch.save(net.state_dict(), 'Model2.pth')

    x_ori = []
    for i in range(len(loss_list)):
        x_ori.append(i)
    plt.title("Graph")
    plt.plot(x_ori, loss_list)
    plt.ylabel("Y")
    plt.xlabel("X")
    plt.show()