cifar目标检测

"""
    使用torch的torchvision加载初始化数据集   --包含了数据集和数据转换工具
    定义卷积神经网络
    定义损失函数
    根据训练数据训练网络
    测试数据上测试网络
"""
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.optim as optim

"""
    dataset  定义好数据的格式和数据变换形式
    dataloader  用iteration 迭代的形式不断读入批次数据
"""

transform = transforms.Compose([
    # 操作打包
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 00.5))
])

# 加载训练集
trainset = torchvision.datasets.CIFAR10(root='./baseLearning/study1', train=True,
                                        download=True, transform=transform)
# 加载测试集
testset = torchvision.datasets.CIFAR10(root='./baseLearning/study1', train=False,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=0)
testloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=False, num_workers=0)

classes = ('airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

def imshow(img):
    """
        恢复为正常图片，通道转化，然后显示
    :return:
    """
    # tensor[batch, channel, H, W]
    # 反归一化，显示图片
    img = img / 2 + 0.5
    # tensor[batch, channel, H, W] --> [channel, H, W]
    npimg = img.numpy()
    # [channel, H, W] -->正常显示图片[H, W, channel]
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

class studyNet(nn.Module):
    """
    自定义自己的网络
    """
    def __init__(self):
        super(studyNet, self).__init__()
        # output = (W - F + 2 * P) / S + 1  W-图片的宽高，F-卷积核大小，P-padding，补边， S-卷积步长
        # 2d卷积 输出通道，输入通道，卷积核
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 输入参数，输出参数
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # tensor[batch, channel, H, W]
        x = self.conv1(x)
        x = F.relu(x)
        # x = F.max_pool2d(x, 2) 都是一样的
        x = F.max_pool2d(x, (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, x.size()[1:].numel())
        x = F.relu(self.fc1(x))
        x = F.relu((self.fc2(x)))
        x = self.fc3(x)
        return x

`在这里插入代码片`dataiter = iter(trainloader)
images, labels = dataiter.next()
# 多张图片拼合一起显示
imshow(torchvision.utils.make_grid(images))
print(len(labels))
# print(labels[0],classes[labels[0]])
print(' '.join(classes[labels[j]] for j in range(len(labels))))

net = studyNet()
print(net)

# 交叉熵损失函数
criterion = nn.CrossEntropyLoss()

# 定义优化器 更新的网络参数，学习率,动量
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(2):

    running_loss = 0.0
    # 0表示从0批次开始读取信息
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        # 梯度归零
        optimizer.zero_grad()
        # 得到预测值
        outputs = net(inputs)
        # 计算损失函数
        loss = criterion(outputs, labels)
        # 后向传播，更新梯度信息，优化模型
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        if i % 2000 == 1999:
            print('[%d, %5d] loss:%.3f' %(epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0
print('Finish')
# 保存训练好的参数
Path = './cifar_net.pth'
torch.save(net.state_dict(), Path)


# 测试图片
dataiter = iter(testloader)
images, labels = dataiter.next()
imshow(torchvision.utils.make_grid(images))
print('GroundTruth:', ' '.join('%5s' % classes[labels[j]] for j in range(len(labels))))

# 加载权重文件
testNet = studyNet()
testNet.load_state_dict(torch.load(Path))
# 输出预测
outputs = testNet(images)
print(outputs)
# 0 表示按列最大值返回，1表示按行最大值返回,_ --->输出最大值,prediction输出
# 最大值的索引
_, predictions = torch.max(outputs, 1)
print('predicted:', ' '.join(classes[predictions[j]] for j in range(len(predictions))))

"""========================================================================="""
correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = testNet(images)
        _, prediction = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (prediction == labels).sum().item()
correct_rate = 100 * (correct / total)
print(correct_rate)