MobileNet-v2 pytorch 代码实现

MobileNet-v2 pytorch 代码实现

标签（空格分隔）： Pytorch 源码

---

主函数

import torch.backends.cudnn as cudnn

from cifar10data import CIFAR10Data
from model import MobileNetV2
from train import Train
from utils import parse_args, create_experiment_dirs


def main():
    # Parse the JSON arguments
    config_args = parse_args()

    # Create the experiment directories
    _, config_args.summary_dir, config_args.checkpoint_dir = create_experiment_dirs(
        config_args.experiment_dir)

    model = MobileNetV2(config_args)#建立图模型

    if config_args.cuda:
        model.cuda()
        cudnn.enabled = True
        cudnn.benchmark = True

    print("Loading Data...")
    data = CIFAR10Data(config_args)#数据流读入
    print("Data loaded successfully\n")


    trainer = Train(model, data.trainloader, data.testloader, config_args)

    if config_args.to_train:
        try:
            print("Training...")
            trainer.train()
            print("Training Finished\n")
        except KeyboardInterrupt:
            pass

    if config_args.to_test:
        print("Testing...")
        trainer.test(data.testloader)
        print("Testing Finished\n")


if __name__ == "__main__":
    main()

model.py

import torch.nn as nn

from layers import inverted_residual_sequence, conv2d_bn_relu6


# 建立网络图模型
class MobileNetV2(nn.Module):
    def __init__(self, args):
        super(MobileNetV2, self).__init__()

        # 配置某些block的stride，满足downsampling的需求
        s1, s2 = 2, 2
        if args.downsampling == 16:
            s1, s2 = 2, 1
        elif args.downsampling == 8:
            s1, s2 = 1, 1

        '''
        network_settings网络的相关配置，从该参数可以看出，Mobile-Net由9个部分组成,
        姑且叫做Mobile block。
        network_settings中:
        't'表示Inverted Residuals的扩征系数
        'c'表示该block输出的通道数
        ‘n’表示当前block由几个残差单元组成
        's'表示当前block的stride
        '''
        # Network is created here, then will be unpacked into nn.sequential
        self.network_settings = [{'t': -1, 'c': 32, 'n': 1, 's': s1},
                                 {'t': 1, 'c': 16, 'n': 1, 's': 1},
                                 {'t': 6, 'c': 24, 'n': 2, 's': s2},
                                 {'t': 6, 'c': 32, 'n': 3, 's': 2},
                                 {'t': 6, 'c': 64, 'n': 4, 's': 2},
                                 {'t': 6, 'c': 96, 'n': 3, 's': 1},
                                 {'t': 6, 'c': 160, 'n': 3, 's': 2},
                                 {'t': 6, 'c': 320, 'n': 1, 's': 1},
                                 {'t': None, 'c': 1280, 'n': 1, 's': 1}]
        self.num_classes = args.num_classes

        ###############################################################################################################

        # Feature Extraction part
        # Layer 0
        # args.width_multiplier网络的通道"瘦身"系数
        # block 0
        self.network = [conv2d_bn_relu6(args.num_channels,
                            int(self.network_settings[0]['c'] * args.width_multiplier),args.kernel_size,
self.network_settings[0]['s'], args.dropout_prob)]

        # Layers from 1 to 7
        for i in range(1, 8):
        # inverted_residual_sequence 根据当前network_settings[i]的配置建立图模型
            self.network.extend(
                inverted_residual_sequence(
                    int(self.network_settings[i - 1]['c'] * args.width_multiplier),
                    int(self.network_settings[i]['c'] * args.width_multiplier),
                    self.network_settings[i]['n'], self.network_settings[i]['t'],
                    args.kernel_size, self.network_settings[i]['s']))

        # Last layer before flattening
        self.network.append(
            conv2d_bn_relu6(int(self.network_settings[7]['c'] * args.width_multiplier), int(self.network_settings[8]['c'] * args.width_multiplier),1 , self.network_settings[8]['s'], args.dropout_prob))

        ###############################################################################################################

        # Classification part
        # 以上输出的特征图进行池化 分类
        self.network.append(nn.Dropout2d(args.dropout_prob, inplace=True))
        self.network.append(nn.AvgPool2d(
            (args.img_height // args.downsampling, args.img_width // args.downsampling)))
        self.network.append(nn.Dropout2d(args.dropout_prob, inplace=True))
        self.network.append(
            nn.Conv2d(int(self.network_settings[8]['c'] * args.width_multiplier), self.num_classes,1, bias=True))

        self.network = nn.Sequential(*self.network)

        self.initialize()

    def forward(self, x): # MobileNetV2的前向传播
        # Debugging mode
        # for op in self.network:
        #     x = op(x)
        #     print(x.shape)
        x = self.network(x)
        x = x.view(-1, self.num_classes)
        return x
    # 初始化权重函数
    def initialize(self):
        """Initializes the model parameters"""
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
                nn.init.xavier_normal(m.weight)
                if m.bias is not None:
                    nn.init.constant(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant(m.weight, 1)
                nn.init.constant(m.bias, 0)

inverted_residual_sequence、InvertedResidualBlock、conv2d_bn_relu6

def inverted_residual_sequence(in_channels, out_channels, num_units, expansion_factor=6,kernel_size=3,initial_stride=2):
    bottleneck_arr = [
        InvertedResidualBlock(in_channels, out_channels, expansion_factor, kernel_size,initial_stride) # 第一个单元stride=initial_stride 后续 stride=1
    for i in range(num_units - 1):
        bottleneck_arr.append(
            InvertedResidualBlock(out_channels, out_channels, expansion_factor, kernel_size, 1))

    return bottleneck_arr

class InvertedResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, expansion_factor=6, kernel_size=3, stride=2):
        super(InvertedResidualBlock, self).__init__()

        if stride != 1 and stride != 2:
            raise ValueError("Stride should be 1 or 2")

        self.block = nn.Sequential(
            nn.Conv2d(in_channels, in_channels * expansion_factor, 1, bias=False), # 扩展通道
            nn.BatchNorm2d(in_channels * expansion_factor),
            nn.ReLU6(inplace=True),

            nn.Conv2d(in_channels * expansion_factor, in_channels * expansion_factor,
                      kernel_size, stride, 1,
                      groups=in_channels * expansion_factor, bias=False), # depth-wise卷积操作
            nn.BatchNorm2d(in_channels * expansion_factor),
            nn.ReLU6(inplace=True),

            nn.Conv2d(in_channels * expansion_factor, out_channels, 1,
                      bias=False), # 恢复输出通道
            nn.BatchNorm2d(out_channels))

        self.is_residual = True if stride == 1 else False # 当该单元的stide = 1 时采用skip connection
        self.is_conv_res = False if in_channels == out_channels else True # 匹配输入 输出通道的一致性

        # Assumption based on previous ResNet papers: If the number of filters doesn't match,
        # there should be a conv1x1 operation.
        if stride == 1 and self.is_conv_res:
            self.conv_res = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False),
                                          nn.BatchNorm2d(out_channels))

    def forward(self, x):# 前向传播
        block = self.block(x)
        if self.is_residual:
            if self.is_conv_res:
                return self.conv_res(x) + block
            return x + block
        return block

‘’‘
该函数分别进行3x3卷积 BN ReLU6操作
’‘’
def conv2d_bn_relu6(in_channels, out_channels, kernel_size=3, stride=2, dropout_prob=0.0):
    # To preserve the equation of padding. (k=1 maps to pad 0, k=3 maps to pad 1, k=5 maps to pad 2, etc.)
    padding = (kernel_size + 1) // 2 - 1
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False),
        nn.BatchNorm2d(out_channels),
        # For efficiency, Dropout is placed before Relu.
        nn.Dropout2d(dropout_prob, inplace=True),
        # Assumption: Relu6 is used everywhere.
        nn.ReLU6(inplace=True)
    )

train.py

import shutil

import torch.nn as nn
import torch.optim
from tensorboardX import SummaryWriter
from torch.autograd import Variable
from torch.optim.rmsprop import RMSprop
from tqdm import tqdm

from utils import AverageTracker


class Train:
    def __init__(self, model, trainloader, valloader, args):
    """
    关键参数说明：
    model:定义的图模型
    trainloader:训练集的输入
    valloader:测试集的输入
    """
        self.model = model
        self.trainloader = trainloader
        self.valloader = valloader
        self.args = args
        self.start_epoch = 0
        self.best_top1 = 0.0

        # Loss function and Optimizer
        self.loss = None
        self.optimizer = None
        self.create_optimization() #定义网络的优化函数及其配置 

        # Model Loading
        self.load_pretrained_model() # 导入预训练模型
        self.load_checkpoint(self.args.resume_from) # 恢复训练模型

        # Tensorboard Writer
        self.summary_writer = SummaryWriter(log_dir=args.summary_dir) # 统计部分变量在训练时的变化情况

    def train(self):
        for cur_epoch in range(self.start_epoch, self.args.num_epochs):

            # Initialize tqdm
            tqdm_batch = tqdm(self.trainloader,
                              desc="Epoch-" + str(cur_epoch) + "-")

            # Learning rate adjustment
            self.adjust_learning_rate(self.optimizer, cur_epoch) # 调整学习率

            # Meters for tracking the average values
            loss, top1, top5 = AverageTracker(), AverageTracker(), AverageTracker()

            # Set the model to be in training mode (for dropout and batchnorm)
            self.model.train()

            for data, target in tqdm_batch:

                if self.args.cuda:
                    data, target = data.cuda(async=self.args.async_loading), target.cuda(async=self.args.async_loading)
                data_var, target_var = Variable(data), Variable(target)

                # Forward pass
                output = self.model(data_var)
                cur_loss = self.loss(output, target_var)

                # Optimization step
                self.optimizer.zero_grad()
                cur_loss.backward()
                self.optimizer.step()

                # Top-1 and Top-5 Accuracy Calculation
                cur_acc1, cur_acc5 = self.compute_accuracy(output.data, target, topk=(1, 5))
                loss.update(cur_loss.data[0])
                top1.update(cur_acc1[0])
                top5.update(cur_acc5[0])

            # Summary Writing
            self.summary_writer.add_scalar("epoch-loss", loss.avg, cur_epoch)
            self.summary_writer.add_scalar("epoch-top-1-acc", top1.avg, cur_epoch)
            self.summary_writer.add_scalar("epoch-top-5-acc", top5.avg, cur_epoch)

            # Print in console
            tqdm_batch.close()
            print("Epoch-" + str(cur_epoch) + " | " + "loss: " + str(
                loss.avg) + " - acc-top1: " + str(
                top1.avg)[:7] + "- acc-top5: " + str(top5.avg)[:7])

            # Evaluate on Validation Set
            if cur_epoch % self.args.test_every == 0 and self.valloader:
                self.test(self.valloader, cur_epoch)

            # Checkpointing
            is_best = top1.avg > self.best_top1
            self.best_top1 = max(top1.avg, self.best_top1)
            self.save_checkpoint({
                'epoch': cur_epoch + 1,
                'state_dict': self.model.state_dict(),
                'best_top1': self.best_top1,
                'optimizer': self.optimizer.state_dict(),
            }, is_best)

    def test(self, testloader, cur_epoch=-1):
        loss, top1, top5 = AverageTracker(), AverageTracker(), AverageTracker()

        # Set the model to be in testing mode (for dropout and batchnorm)
        self.model.eval()

        for data, target in testloader:
            if self.args.cuda:
                data, target = data.cuda(async=self.args.async_loading), target.cuda(
                    async=self.args.async_loading)
            data_var, target_var = Variable(data, volatile=True), Variable(target, volatile=True)

            # Forward pass
            output = self.model(data_var)
            cur_loss = self.loss(output, target_var)

            # Top-1 and Top-5 Accuracy Calculation
            cur_acc1, cur_acc5 = self.compute_accuracy(output.data, target, topk=(1, 5))
            loss.update(cur_loss.data[0])
            top1.update(cur_acc1[0])
            top5.update(cur_acc5[0])

        if cur_epoch != -1:
            # Summary Writing
            self.summary_writer.add_scalar("test-loss", loss.avg, cur_epoch)
            self.summary_writer.add_scalar("test-top-1-acc", top1.avg, cur_epoch)
            self.summary_writer.add_scalar("test-top-5-acc", top5.avg, cur_epoch)

        print("Test Results" + " | " + "loss: " + str(loss.avg) + " - acc-top1: " + str(
            top1.avg)[:7] + "- acc-top5: " + str(top5.avg)[:7])

    def save_checkpoint(self, state, is_best, filename='checkpoint.pth.tar'):
        torch.save(state, self.args.checkpoint_dir + filename)
        if is_best:
            shutil.copyfile(self.args.checkpoint_dir + filename,
                            self.args.checkpoint_dir + 'model_best.pth.tar')

    def compute_accuracy(self, output, target, topk=(1,)):
        """Computes the accuracy@k for the specified values of k"""
        maxk = max(topk)
        batch_size = target.size(0)

        _, idx = output.topk(maxk, 1, True, True)
        idx = idx.t()
        correct = idx.eq(target.view(1, -1).expand_as(idx))

        acc_arr = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            acc_arr.append(correct_k.mul_(1.0 / batch_size))
        return acc_arr

    def adjust_learning_rate(self, optimizer, epoch):
        """Sets the learning rate to the initial LR multiplied by 0.98 every epoch"""
        learning_rate = self.args.learning_rate * (self.args.learning_rate_decay ** epoch)
        for param_group in optimizer.param_groups:
            param_group['lr'] = learning_rate

    def create_optimization(self):
        self.loss = nn.CrossEntropyLoss()

        if self.args.cuda:
            self.loss.cuda()

        self.optimizer = RMSprop(self.model.parameters(), self.args.learning_rate,
                                 momentum=self.args.momentum,
                                 weight_decay=self.args.weight_decay)

    def load_pretrained_model(self):
        try:
            print("Loading ImageNet pretrained weights...")
            pretrained_dict = torch.load(self.args.pretrained_path)
            self.model.load_state_dict(pretrained_dict)
            print("ImageNet pretrained weights loaded successfully.\n")
        except:
            print("No ImageNet pretrained weights exist. Skipping...\n")

    def load_checkpoint(self, filename):
        filename = self.args.checkpoint_dir + filename
        try:
            print("Loading checkpoint '{}'".format(filename))
            checkpoint = torch.load(filename)
            self.start_epoch = checkpoint['epoch']
            self.best_top1 = checkpoint['best_top1']
            self.model.load_state_dict(checkpoint['state_dict'])
            self.optimizer.load_state_dict(checkpoint['optimizer'])
            print("Checkpoint loaded successfully from '{}' at (epoch {})\n"
                  .format(self.args.checkpoint_dir, checkpoint['epoch']))
        except:
            print("No checkpoint exists from '{}'. Skipping...\n".format(self.args.checkpoint_dir))

总结

Mobile-Net v2 netscope : http://ethereon.github.io/netscope/#/gist/d01b5b8783b4582a42fe07bd46243986

Pytorch与caffe实现的Mobile-Net v2 不完全一样
1. 在caffe中，每一个block的前面几个单元采用跳跃连接，stride = 1，最后一个单元不采用跳跃链接，stride = 2
2. 在pytorch中，相反。

现在来说说Mobile-Net v2的特点：

1. 在正常的残差单元，depth_bottleneck = inchannel / 4，而inverted residual unit的 depth_bottleneck = inchannel * 6，具体件论文；
2. 在通道减小的卷积层中不采用非线性激活函数（ReLU6）