Pytorch實現cifar-10圖像分類

導入必要的包

import torch
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

通過transform 實現對數據進行處理

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ])

數據加載：

trainset = torchvision.datasets.CIFAR10(root=r'./data',
                                        train=True,download=False,transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size= 4, shuffle=False, num_workers=0)

testset = torchvision.datasets.CIFAR10(root=r'./data',
                                       train=False,download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=0)

classes = ('plane', 'car', 'bird','cat','deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

cifar-10數據可以事先下載好，放到 data文件夾下。 若網速夠快，可以在線下載：

設置： download=True 

展示一下我們加載的cifar-10數據：

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

dataiter = iter(trainloader)
images, labels = dataiter.next()
imshow(torchvision.utils.make_grid(images))

print(' '.join('%5s' %classes[labels[j]] for j in range(4)))

定義完了模型，這裏只是設置了2層卷積+2層池化+3層全連接的網絡結構

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(6, 16, 5)

        self.fc1 = nn.Linear(16*5*5, 120)   # 6*16*5 全連接線性化
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))

        x = x.view(-1, 16*5*5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

設置損失函數和優化器：

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

定義訓練輪次，進行訓練：

epochs = 5
for epoch in range(epochs):
    running_loss = 0.0
    for i , data in enumerate(trainloader, 0):
        inputs, labels = data
        # zero the parameter gradients
        optimizer.zero_grad()

        # forward +backward +optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if i%2000==1999:    # print every 2000 mini-batches
            print('%d, %5d, loss: % .3f' %(epoch+1, i+1, running_loss/2000))
            running_loss = 0 # make loss zeros

保存訓練好的模型：

# save model
torch.save(net.state_dict(), './save_model') # only save weights

加載模型：

#load model
model = Net()
model.load_state_dict(torch.load('./save_model'))

進行預測：

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

對整個測試集進行預測，計算準確率

# test the correct
correct = 0
total = 0
with torch.no_grad():   # 實現一定速度的提升，並節省一半的顯存，因爲其不需要保存梯度
    for data in testloader:
        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))
# Accuracy of the network on the 10000 test images: 60%

使用GPU、或多個GPU進行網絡訓練時，設置：

# 如何在GPU 上訓練
device= torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)   # could look up cuda

model.to(device)
inputs, labels = inputs.to(device), labels.to(device)

# 用 DataParallel 使用多個GPU
# Pytorch 默認只會使用一個GPU，可以通過DataParallel 讓你的模型並行運行
if torch.cuda.device_count()>1:
    print("Let's use ", torch.cuda.device_count(), "GPUs" )
    model = nn.DataParallel(model)
model.to(device)