通过pytorch搭建mobilenet_v2网络

import torch
from torch import nn
import torchvision
from torchvision import transforms
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader

train_dir = './training/training'
valid_dir = './validation/validation'
train_transformation = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[1, 1, 1])
])
valid_transformation = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[1, 1, 1])
])
train_ds = torchvision.datasets.ImageFolder(train_dir, train_transformation)
valid_ds = torchvision.datasets.ImageFolder(valid_dir, valid_transformation)
train_dl = DataLoader(train_ds, batch_size=32, shuffle=True)
valid_dl = DataLoader(valid_ds, batch_size=32)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class ConvBnRelu6(nn.Module):
    def __init__(self, in_channel, out_channel, kernel_size, strides, padding):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channel, out_channel, kernel_size, strides, padding, bias=False),
            nn.BatchNorm2d(out_channel),
            nn.ReLU()
        )

    def forward(self, inputs):
        x = self.conv(inputs)
        return x


class DWBnRelu6(nn.Module):
    def __init__(self, in_channel, out_channel, kernel_size, strides, padding):
        super().__init__()
        self.dw = nn.Sequential(
            nn.Conv2d(in_channel, out_channel, kernel_size, strides, padding, groups=in_channel),
            nn.BatchNorm2d(in_channel),
            nn.ReLU6()
        )

    def forward(self, inputs):
        x = self.dw(inputs)
        return x


class BottleNeckDown(nn.Module):
    def __init__(self, in_channel, out_channel, t):
        super().__init__()
        exp_channel = in_channel * t
        self.branch1 = nn.Sequential(
            ConvBnRelu6(in_channel, exp_channel, 1, 1, 0),
            DWBnRelu6(exp_channel, exp_channel, 3, 2, 1),
            nn.Conv2d(exp_channel, out_channel, 1, 1, 0),
            nn.BatchNorm2d(out_channel)
        )

    def forward(self, inputs):
        x = self.branch1(inputs)
        return x


class BottleNeck(nn.Module):
    def __init__(self, in_channel, out_channel, t):
        super().__init__()
        self.in_channel = in_channel
        self.out_channel = out_channel
        exp_channel = in_channel * t
        self.branch1 = nn.Sequential(
            ConvBnRelu6(in_channel, exp_channel, 1, 1, 0),
            DWBnRelu6(exp_channel, exp_channel, 3, 1, 1),
            nn.Conv2d(exp_channel, out_channel, 1, 1, 0),
            nn.BatchNorm2d(out_channel)
        )

    def forward(self, inputs):
        branch = self.branch1(inputs)
        if self.in_channel == self.out_channel:
            branch2 = inputs
            branch = branch + branch2
        return branch


class MobileNetV2(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.head = nn.Sequential(
            ConvBnRelu6(3, 32, 3, 2, 1)
        )
        net_shape = [(1, 16, 1, 1, 32),
                     (6, 24, 2, 2, 16),
                     (6, 32, 3, 2, 24),
                     (6, 64, 4, 2, 32),
                     (6, 96, 3, 1, 64),
                     (6, 160, 3, 2, 96),
                     (6, 320, 1, 1, 160)]
        self.body = nn.Sequential()
        for t, c, n, s, in_channel in net_shape:
            if s == 1:
                self.body.append(BottleNeck(in_channel, c, t))
            else:
                for i in range(n):
                    if i == 0:
                        self.body.append(BottleNeckDown(in_channel, c, t))
                    else:
                        self.body.append(BottleNeck(c, c, t))
        self.top = nn.Sequential(
            ConvBnRelu6(320, 1280, 1, 1, 0),
            nn.MaxPool2d(7),
            nn.Conv2d(1280, num_classes, 1, 1, 0)
        )
        self.activation = nn.Softmax()

    def forward(self, inputs):
        x = self.head(inputs)
        x = self.body(x)
        x = self.top(x)
        x = nn.Flatten()(x)
        x = self.activation(x)
        return x


def cal_acc(y_pred, y_true):
    acc = (torch.argmax(y_pred, dim=1) == y_true).sum().item()
    return acc


def train_step(model, epoch, train_dl, valid_dl, loss_func, optimizer, lr_schedule):
    total_train_loss = 0.0
    total_train_acc = 0.0
    total_train = 0
    total_valid_loss = 0.0
    total_valid_acc = 0.0
    total_valid = 0
    model.train()
    for x, y in train_dl:
        x, y = x.to(device), y.to(device)
        prediction = model(x)
        train_loss = loss_func(prediction, y)
        optimizer.zero_grad()
        train_loss.backward()
        optimizer.step()
        with torch.no_grad():
            total_train_loss += train_loss.item()
            total_train += y.size(0)
            train_acc = cal_acc(prediction, y)
            total_train_acc += train_acc
            print(
                '\repoch: %d, train_loss: %3.3f, train_acc: %3.3f' % (epoch, train_loss.item(), train_acc / y.size(0)),
                end='')
    total_train_loss = total_train_loss / total_train
    total_train_acc = total_train_acc / total_train
    lr_schedule.step()
    model.eval()
    with torch.no_grad():
        for x, y in valid_dl:
            x, y = x.to(device), y.to(device)
            prediction = model(x)
            total_valid_loss += loss_func(prediction, y).item()
            total_valid_acc += cal_acc(prediction, y)
            total_valid += y.size(0)
    total_valid_loss = total_valid_loss / total_valid
    total_valid_acc = total_valid_acc / total_valid
    print('epoch: %d, train_loss: %3.3f, train_acc: %3.3f, valid_loss: %3.3f, valid_acc: %3.3f' %
          (epoch, total_train_loss, total_train_acc, total_valid_loss, total_valid_acc))
    return total_train_loss, total_train_acc, total_valid_loss, total_valid_acc


def model_fit(model, train_dl, valid_dl, loss_func, optimizer, lr_schedule, epochs=10):
    history = {'train_loss': [],
               'train_acc': [],
               'valid_loss': [],
               'valid_acc': []}
    for epoch in range(epochs):
        train_loss, train_acc, valid_loss, valid_acc = train_step(model, epoch, train_dl, valid_dl, loss_func,
                                                                  optimizer, lr_schedule)
        history['train_loss'].append(train_loss)
        history['train_acc'].append(train_acc)
        history['valid_loss'].append(valid_loss)
        history['valid_acc'].append(valid_acc)
    return history


def train_model(model, train_dl, valid_dl, epochs=10):
    model = model.to(device)
    loss_fn = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
    lr_schedule = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.95)
    history = model_fit(model, train_dl, valid_dl, loss_fn, optimizer, lr_schedule)
    return history, model


mobile_net = MobileNetV2(10)
history, mobile_net = train_model(mobile_net, train_dl, valid_dl)