MobileNet 实现对CIFAR-10数据集训练(pytorch)

import torch

import torchvision

from torchvision import datasets, transforms

from torch.utils.data import DataLoader

from torch import nn

import torch.nn.functional as F

import numpy as np

from torch.utils.data import Dataset, DataLoader, TensorDataset

 

train_dataset =   datasets.CIFAR10('./cifar10_data/', train=True, download=True,

                       transform=transforms.Compose([

                           #transforms.ToPILImage(),

                           #transforms.Grayscale(),

                           transforms.ToTensor()

                       ]))

 

train_loader = DataLoader(dataset = train_dataset, batch_size = 500, shuffle = True)

 

class Block(nn.Module):

    '''Depthwise conv + Pointwise conv'''

    def __init__(self, in_planes, out_planes, stride=1):

        super(Block, self).__init__()

        self.conv1 = nn.Conv2d\

            (in_planes, in_planes, kernel_size=3, stride=stride, 

             padding=1, groups=in_planes, bias=False)

        self.bn1 = nn.BatchNorm2d(in_planes)

        self.conv2 = nn.Conv2d\

            (in_planes, out_planes, kernel_size=1, 

            stride=1, padding=0, bias=False)

        self.bn2 = nn.BatchNorm2d(out_planes)

 

    def forward(self, x):

        out = F.relu(self.bn1(self.conv1(x)))

        out = F.relu(self.bn2(self.conv2(out)))

        return out

 

class MobileNet(nn.Module):

    # (128,2) means conv planes=128, conv stride=2, 

    # by default conv stride=1

    cfg = [64, (128,2), 128, (256,2), 256, (512,2), 

           512, 512, 512, 512, 512, (1024,2), 1024]

 

    def __init__(self, num_classes=10):

        super(MobileNet, self).__init__()

        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, 

            stride=1, padding=1, bias=False)

        self.bn1 = nn.BatchNorm2d(32)

        self.layers = self._make_layers(in_planes=32)

        self.linear = nn.Linear(1024, num_classes)

 

    def _make_layers(self, in_planes):

        layers = []

        for x in self.cfg:

            out_planes = x if isinstance(x, int) else x[0]

            stride = 1 if isinstance(x, int) else x[1]

            layers.append(Block(in_planes, out_planes, stride))

            in_planes = out_planes

        return nn.Sequential(*layers)

 

    def forward(self, x):

        out = F.relu(self.bn1(self.conv1(x)))

        out = self.layers(out)

        out = F.avg_pool2d(out, 2)

        out = out.view(out.size(0), -1)

        out = self.linear(out)

        return out

 

epoch = 5

 

net = MobileNet()

cost = torch.nn.CrossEntropyLoss()

optimizer = torch.optim.Adam(net.parameters(), lr=0.0005)

 

for k in range(epoch):

    sum_loss = 0.0

    train_correct = 0

    for i, data in enumerate(train_loader, 0):

        inputs, labels = data

        optimizer.zero_grad()

        outputs = net(inputs)

 

        loss = cost(outputs, labels)

        loss.backward()

        optimizer.step()

 

        print(loss)

        _, id = torch.max(outputs.data, 1) 

        sum_loss += loss.data

        train_correct += torch.sum(id == labels.data)

        #print('[%d,%d] loss:%.03f' % (k + 1, k, sum_loss / len(train_loader)))

    print('        correct:%.03f%%' % (100 * train_correct / len(train_dataset)))