pytorch入门（五）—— 卷积神经网络

经典卷积神经网络（CNN）

import torch as tr
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
import torchvision as tv
import matplotlib.pyplot as plt

EPOCH = 3
BATCH_SIZE = 50
LR = 0.001
DOWNLOAD_MNIST = True

train_data = tv.datasets.MNIST(
        root='./data/',
        train=True,
        transform=tv.transforms.ToTensor(),
        download=DOWNLOAD_MNIST
        )
print(train_data.train_data.size())
#tr.Size([60000, 28, 28])
print(train_data.train_labels.size())

#for i in range(1,4):
#    plt.imshow(train_data.train_data[i].numpy(), cmap='gray')
#    plt.title('%i' % train_data.train_labels[i])
#    plt.show()

train_loader = Data.DataLoader(dataset=train_data, 
        batch_size=BATCH_SIZE,
        shuffle=True)

test_data = tv.datasets.MNIST(
        root='./data',
        train=False
        )

test_x = Variable(tr.unsqueeze(test_data.test_data, dim=1)).type(tr.FloatTensor).cuda()
test_y = test_data.test_labels.cuda()

print(test_x.requires_grad)
print(test_y.size(0))

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
                nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=2)
                )

        self.conv2 = nn.Sequential(
                nn.Conv2d(16, 32, 5, 1, 2),
                nn.ReLU(),
                nn.MaxPool2d(2)
                )

        self.out = nn.Linear(32*7*7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        output = self.out(x)
        return output

cnn = CNN()
cnn.cuda()
#print(cnn)

#params = list(cnn.parameters())
#print(params)
#print(len(params))
#print(params[0].size())

optimizer = tr.optim.Adam(cnn.parameters(), lr=LR)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):
        b_x = Variable(x).cuda()
        b_y = Variable(y).cuda()

        output = cnn(b_x)
        loss = loss_fn(output, b_y)
        optimizer.zero_grad()

        loss.backward()
        optimizer.step()

        if step % 100 == 0:
            test_output = cnn(test_x)
            pred_y = tr.max(test_output, 1)[1].data.squeeze()
            accuracy = (pred_y == test_y).sum().float() / test_y.size(0)
            print('Epoch:', epoch, '|Step:', step, '|train loss:%.4f' % loss.item(),
                    '|test accuracy:%.4f' % accuracy)
test_output = cnn(test_x[:20])
pred_y = tr.max(test_output, 1)[1].cpu().data.numpy().squeeze()
print(pred_y[:20], 'prediction number')
print(test_y[:20].cpu().numpy(), 'real number')

深度残差网络模型（ResNet）

import torch as tr
import torch.nn as nn
from torch.autograd import Variable
import torchvision as tv
import torch.utils.data as Data
import torchvision.transforms as tf
import torchvision.datasets as dsets
import matplotlib.pyplot as plt


transform = tf.Compose([
    tf.Resize(40),
    tf.RandomHorizontalFlip(),
    tf.RandomCrop(32),
    tf.ToTensor()
    ])

train_dataset = dsets.CIFAR10(root='./data/',
        train=True,
        transform=transform,
        download=True)

test_dataset = dsets.CIFAR10(root='./data/',
        train=False,
        transform=tf.ToTensor()
        )

train_loader = Data.DataLoader(dataset=train_dataset,
        batch_size=100,
        shuffle=False)

test_loader = Data.DataLoader(dataset=test_dataset,
        batch_size=100,
        shuffle=False)

def conv3x3(in_channels, out_channels, stride=1):
    return nn.Conv2d(in_channels, out_channels, kernel_size=3, 
            stride=stride, padding=1, bias=False)
    #nn.conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1)


class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(ResidualBlock, self).__init__()
        self.conv1 = conv3x3(in_channels, out_channels, stride)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(out_channels, out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = downsample

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.downsample:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)
        return out


class ResNet(nn.Module):
    def __init__(self, block, layers, num_classes=10):
        super(ResNet, self).__init__()
        self.in_channels = 16
        self.conv = conv3x3(3, 16)
        self.bn = nn.BatchNorm2d(16)
        self.relu = nn.ReLU(inplace=True)
        self.layer1 = self.make_layer(block, 16, layers[0])
        self.layer2 = self.make_layer(block, 32, layers[0], 2)
        self.layer3 = self.make_layer(block, 64, layers[1], 2)
        self.avg_pool = nn.AvgPool2d(8)
        self.fc = nn.Linear(64, num_classes)

    def make_layer(self, block, out_channels, blocks, stride=1):
        downsample = None
        if(stride != 1) or (self.in_channels != out_channels):
            downsample = nn.Sequential(
                    conv3x3(self.in_channels, out_channels, stride=stride),
                    nn.BatchNorm2d(out_channels)
                    )
        layers = []
        layers.append(block(self.in_channels, out_channels, stride, downsample))
        self.in_channels = out_channels
        for i in range(1, blocks):
            layers.append(block(out_channels, out_channels))

        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv(x)
        out = self.bn(out)
        out = self.relu(out)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.avg_pool(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out


resnet = ResNet(ResidualBlock, [3, 3, 3])

resnet.cuda()

#print(resnet)

loss_fn = nn.CrossEntropyLoss()
lr = 0.001
optimizer = tr.optim.Adam(resnet.parameters(), lr=lr)

for epoch in range(80):
    for i, (images, labels) in enumerate(train_loader):
        images = Variable(images).cuda()
        labels = Variable(labels).cuda()

        optimizer.zero_grad()
        outputs = resnet(images)
        loss = loss_fn(outputs, labels)
        loss.backward()
        optimizer.step()

        if (i + 1) % 100 == 0:
            print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f' % 
                    (epoch + 1, 80, i + 1, 500, loss.item()))

    if (epoch + 1) % 20 == 0:
        lr /= 3
        optimizer = tr.optim.Adam(resnet.parameters(), lr=lr)


correct = 0
total = 0
for images, labels in test_loader:
    images = Variable(images.cuda())
    outputs = resnet(images)
    _, predicted = tr.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted.cpu() == labels).sum()

print('Accuracy of the model on the test images: %d %%' % (100 * correct / total))

tr.save(resnet, 'resnet.pkl')