基于TensorFlow微调AlexNet

【翻译】基于TensorFlow微调AlexNet

（以下均以原博主视角称呼，基本意译，不好的地方还请见谅）

我已更新了使用新数据输入方式的代码仓库。这个新特性来自于版本1.12rc0以上的TensorFlow，可以阅读我的其他博客来进行了解。以下博文里的链接也是指向这里解释的代码。

---

经过一年多，我终于找到时间和空闲来写我新的博文。这次博文是有关基于TensorFlow 1.0版本微调AlexNet。你可能会说这AlexNet是2012年的网络了，现在都是2017年了！那我就先讲讲我为什么认为这个网络值得我们花力气去微调的原因。

• 尽管有很多的网络模型微调的教程，但是大部分都是有关于VGG模型或者是Inception模型，并没有AlexNet模型。尽管网络模型微调背后的思想是一样的，但是最大的不同在于TensorFlow（Keras也是）已经含有VGG/Inception类以及相关参数（基于ImageNet数据集预训练得到）。对于AlexNet模型，我们需要自己做一点微小的工作。
• 另一个原因是，在我大多数的个人项目中，AlexNet模型的效果就相当好了，就没什么理由为了可能提高的0.5%准确度去换一个有着巨量参数的庞大模型。当模型越来越深的时候，它们的计算就会越来越费时间。这在一些项目中是无法接受的。
• 过去，我主要使用Caffe来微调卷积网络。但是说实话，我发现使用Caffe相当繁琐，例如需要通过.prototxt定义模型、一团的模型结构等等。使用了Keras一段时间后，我最近转移到了TensorFlow上来。现在我想要像以前一样能够在TensorFlow上微调同样的网络模型。
• 当然，我认为写这个博文的主要原因就是做一个像这样的项目，能够帮助我更好的理解TensorFlow。

无责任声明

After finishing to write this article I ended up having written another very long post(译者翻译不好这句话). 以下基本上分为两个部分：
在第一部分中，我创建了一个类来定义AlexNet模型图，并带有加载预训练参数的函数；
第二部分中讲了针对新的数据集，如何使用这个类来微调AlexNet得到所需要的模型。虽然我推荐阅读第一部分，但是想直接知道如何微调AlexNet可以点击这里跳过第一部分。

另外，本博文并不是TensorFlow的教程也适用于一般卷积网络的微调。我会解释大部分你们需要做的工作，但是你们要懂得关于TensorFlow和机器/深度学习的基本知识来理解文中的东西。

获取预训练权值

不像VGG或Inception那样，TensorFlow中没有预训练好的AlexNet模型。Caffe有，但是将其手动转化到TensorFlow可用的结构实在繁琐。幸运的是，有一个可以用于转化的小工具，可以将任何*prototxt模型定义从caffe转化为python代码和一个TensorFlow模型，同样适用权值。我自己尝试了一下，发现这个效果非常好。不管怎样，你可以在这里下载到已经转换好的参数文件(bvlc_alexnet.npy)。

模型结构

在开始微调AlexNet之前，我们必须先创建所谓的”Graph of the Model”。这和我在上个博客中为BatchNormalization定义的一样，只是这里是针对整个模型。不过不用担心，我们不必亲自做所有的事情。首先让我们先看看原论文中展示模型结构：

Architecture of AlexNet, as shown in the original paper (link above).

值得注意的是一些卷积层分为了并列的两个（层二、层四和层五）。这是为了将卷积层的计算分配至两个GPU进行计算（我猜是因为那时候的GPU还不够强大。译者：对，你猜对了，+10分）。尽管如今这可能没有必要，但是为了重现AlexNet的结果，我们还是要定义同样的分割，只不过我们只使用一个GPU。

所以我们开始吧：我们为模型定义一个有如下结构的类（整个代码能够在Github上找到->这里）。Note:查看代码新版本的更新信息。

class AlexNet(object):

    def __init__(self, x, keep_prob, num_classes, skip_layer, weights_path = 'DEFAULT'):
        """
        Inputs:
        - x: tf.placeholder, for the input images
        - keep_prob: tf.placeholder, for the dropout rate
        - num_classes: int, number of classes of the new dataset
        - skip_layer: list of strings, names of the layers you want to reinitialize
        - weights_path: path string, path to the pretrained weights, 
          (if bvlc_alexnet.npy is not in the same folder)
        """
        # Parse input arguments
        self.X = x
        self.NUM_CLASSES = num_classes
        self.KEEP_PROB = keep_prob
        self.SKIP_LAYER = skip_layer
        self.IS_TRAINING = is_training
        if weights_path == 'DEFAULT':
            self.WEIGHTS_PATH = 'bvlc_alexnet.npy'
        else:
            self.WEIGHTS_PATH = weights_path
        # Call the create function to build the computational graph of AlexNet
        self.create()

    def create(self):
        pass

    def load_initial_weights(self):
        pass

在__init__函数中，我们将输入元解析为类变量并调用create函数。我们本能够一次写完所有的代码，但是我个人认为这样看起来更清楚。load_initial_weights函数用来将预训练的参数赋给我们创建的变量。

辅助函数

现在我们有了类的基本结构，接下来我们来定义一些用于创建网络层的辅助函数。由于我们要在一个卷积层函数里实现分开的功能与不分开的功能，这个函数可能会‘麻烦’。下面这个函数是caffe-to-tensorflow repo中采用的版本：

def conv(x, filter_height, filter_width, num_filters, stride_y, stride_x, name,padding='SAME', groups=1):
    # Get number of input channels
    input_channels = int(x.get_shape()[-1])

    # Create lambda function for the convolution
    convolve = lambda i, k: tf.nn.conv2d(i, k,
                                        strides = [1, stride_y, stride_x, 1],
                                        padding = padding)

    with tf.variable_scope(name) as scope:
        # Create tf variables for the weights and biases of the conv layer
        weights = tf.get_variable('weights',
                                shape = [filter_height, filter_width,
                                input_channels/groups, num_filters])
        biases = tf.get_variable('biases', shape = [num_filters])

        if groups == 1:
            conv = convolve(x, weights)

        # In the cases of multiple groups, split inputs & weights and
        else:
            # Split input and weights and convolve them separately
            input_groups = tf.split(axis = 3, num_or_size_splits=groups, value=x)
            weight_groups = tf.split(axis = 3, num_or_size_splits=groups, value=weights)
            output_groups = [convolve(i, k) for i,k in zip(input_groups, weight_groups)]
            # Concat the convolved output together again
            conv = tf.concat(axis = 3, values = output_groups)

        # Add biases
        bias = tf.reshape(tf.nn.bias_add(conv, biases), conv.get_shape().as_list())

        # Apply relu function
        relu = tf.nn.relu(bias, name = scope.name)

        return relu

在一个函数中使用lambda匿名函数和递推式构造列表来处理是一个相当简洁的方式。我希望我在其余的代码中写的注释足够明白了。接下来是定义全连接层的函数。这个就很简单了：

def fc(x, num_in, num_out, name, relu = True):
    with tf.variable_scope(name) as scope:

    # Create tf variables for the weights and biases
    weights = tf.get_variable('weights', shape=[num_in, num_out], trainable=True)
    biases = tf.get_variable('biases', [num_out], trainable=True)

    # Matrix multiply weights and inputs and add bias
    act = tf.nn.xw_plus_b(x, weights, biases, name=scope.name)

    if relu == True:
        # Apply ReLu non linearity
        relu = tf.nn.relu(act)
        return relu
    else:
        return act

Note:我知道这个函数可以用行数更少的代码实现(如tf.nn.relu_layer())，但是像这样就允许我们将激活值加入到tf.summary()里，然后在TensorBoard中来监控训练中的激活值。余下的是最大池化层、LRN层(Local-Response-Normalization)和Dropout层，并且应该也很明白的注释了：

def max_pool(x, filter_height, filter_width, stride_y, stride_x, name, padding='SAME'):
    return tf.nn.max_pool(x, ksize=[1, filter_height, filter_width, 1],
                             strides = [1, stride_y, stride_x, 1],
                             padding = padding, name = name)

def lrn(x, radius, alpha, beta, name, bias=1.0):
    return tf.nn.local_response_normalization(x, depth_radius = radius,
                                                 alpha = alpha, beta = beta,
                                                 bias = bias, name = name)

def dropout(x, keep_prob):
    return tf.nn.dropout(x, keep_prob)

建立AlexNet图

现在我们把create函数完善了，建立模型图。

def create(self):

    # 1st Layer: Conv (w ReLu) -> Pool -> Lrn
    conv1 = conv(self.X, 11, 11, 96, 4, 4, padding = 'VALID', name = 'conv1')
    pool1 = max_pool(conv1, 3, 3, 2, 2, padding = 'VALID', name = 'pool1')
    norm1 = lrn(pool1, 2, 2e-05, 0.75, name = 'norm1')

    # 2nd Layer: Conv (w ReLu) -> Pool -> Lrn with 2 groups
    conv2 = conv(norm1, 5, 5, 256, 1, 1, groups = 2, name = 'conv2')
    pool2 = max_pool(conv2, 3, 3, 2, 2, padding = 'VALID', name ='pool2')
    norm2 = lrn(pool2, 2, 2e-05, 0.75, name = 'norm2')

    # 3rd Layer: Conv (w ReLu)
    conv3 = conv(norm2, 3, 3, 384, 1, 1, name = 'conv3')

    # 4th Layer: Conv (w ReLu) splitted into two groups
    conv4 = conv(conv3, 3, 3, 384, 1, 1, groups = 2, name = 'conv4')

    # 5th Layer: Conv (w ReLu) -> Pool splitted into two groups
    conv5 = conv(conv4, 3, 3, 256, 1, 1, groups = 2, name = 'conv5')
    pool5 = max_pool(conv5, 3, 3, 2, 2, padding = 'VALID', name = 'pool5')

    # 6th Layer: Flatten -> FC (w ReLu) -> Dropout
    flattened = tf.reshape(pool5, [-1, 6*6*256])
    fc6 = fc(flattened, 6*6*256, 4096, name='fc6')
    dropout6 = dropout(fc6, self.KEEP_PROB)

    # 7th Layer: FC (w ReLu) -> Dropout
    fc7 = fc(dropout6, 4096, 4096, name = 'fc7')
    dropout7 = dropout(fc7, self.KEEP_PROB)

    # 8th Layer: FC and return unscaled activations
    # (for tf.nn.softmax_cross_entropy_with_logits)
    self.fc8 = fc(dropout7, 4096, self.NUM_CLASSES, relu = False, name='fc8')

注意，在定义最后一层的时候，我们使用了self.NUM_CLASSES变量，因此我们能够使用这个类和函数解决不同的分类问题。看，这就是我们要的图，而且这个看起来多酷，太让我喜欢了。下面是TensorBoard画的整个网络的计算图。

Visualization of the computational graph of Tensorboard (left) and a closer look to the conv5 layer (right), one of the layers with splitting.

加载预训练参数

好了，现在到了load_initial_weights函数。这个函数的目的将存储于self.WEIGHTS_PATH的预训练参数赋值给那些没有在self.SKIP_LAYER中指定的网络层的参数，这些网络层是我们想要进行训练的。打开bvlc_alexnet.npy文件观察一下参数的结构，你就会注意到它们是以链表为值的Python字典。一个键对应一层，并包含有权重、偏置值的列表。如果你自己用caffe-to-tensorflow函数去转换参数，你会得到值为字典的Python字典，如weights['conv1']是一个键为weights和biases字典。接下来我要来写参数（参数文件可以在这下载）的函数了。对于每个参数列表项，我们必须先确认它的维度然后才能决定它赋给权重weights(len(shape)>1)还是偏置biases(len(shape)==1)。

def load_initial_weights(self, session):
    # Load the weights into memory
    weights_dict = np.load(self.WEIGHTS_PATH, encoding = 'bytes').item()

    # Loop over all layer names stored in the weights dict
    for op_name in weights_dict:

        # Check if the layer is one of the layers that should be reinitialized
        if op_name not in self.SKIP_LAYER:
            with tf.variable_scope(op_name, reuse = True):
                # Loop over list of weights/biases and assign them to their corresponding tf variable
                for data in weights_dict[op_name]:
                    # Biases
                    if len(data.shape) == 1:
                        var = tf.get_variable('biases', trainable = False)
                        session.run(var.assign(data))

                    # Weights
                    else:
                        var = tf.get_variable('weights', trainable = False)
                        session.run(var.assign(data))

写完这段代码，整个AlexNet类就完成了！

测试

为了测试这个模型是不是正确和参数是否被正确赋值，我们可以创建一个ImageNet原始模型（最后一层有1000个类别）并将预训练参数赋给所有网络层。我从原始ImageNet数据库里抓取了几张图片进行预测分类，下面是分类结果：

Three images taken from the ImageNet Database and tested with the implemented AlexNet class. IPython notebook to reproduce the results can be found here.

结果看上去很好，所以接下来我们进入微调部分。

微调AlexNet

经过前面那么长的篇幅，终于到了本文最核心的地方：使用创建好的AlexNet类进行微调网络来适用于你自己的数据。这个主要思想现在相当明确：我们会创建一个模型，通过在skip_layer变量中记录网络层的名字来跳过最后几层网络，然后在TensorFlow里建立损失函数和优化操作，最后开始一个会话Session训练网络。在TensorBoard的支持下，我们能够通过一些操作观察到训练过程。为了测试微调的程序，我从Kaggle猫狗大战竞赛里下载train.zip文件

Image from the kaggle competition

我进一步将图片集分为了训练集、验证集和测试集（70-15-15），并为每个数据集创建了一个文本文件（.txt），这个文件里包含了图片的路径和类别标签。有了这个文本文件后，我又创建了另一个类作为图片数据生成器（就像Keras里的例子）。我知道又更好的方式，但是之前有另一个项目需要我考虑如何加载和预处理图片，然后已经有了这个脚本（译者：程序员的‘懒’），所以我就简单地赋值到这里来了。这段代码可以在github仓库里找到。另外，我个人出于教育目的喜欢用更多的脚本，不会把代码写成可调用的函数而是写成大家可以打开、阅读的脚本，这样能更好地明白发生了什么。

构造部分

模块导入之后，我首先定义所有需要的变量。

import numpy as np
import tensorflow as tf
from datetime import datetime
from alexnet import AlexNet
from datagenerator import ImageDataGenerator

# Path to the textfiles for the trainings and validation set
train_file = '/path/to/train.txt'
val_file = '/path/to/val.txt'

# Learning params
learning_rate = 0.001
num_epochs = 10
batch_size = 128

# Network params
dropout_rate = 0.5
num_classes = 2
train_layers = ['fc8', 'fc7']

# How often we want to write the tf.summary data to disk
display_step = 1

# Path for tf.summary.FileWriter and to store model checkpoints
filewriter_path = "/tmp/finetune_alexnet/dogs_vs_cats"
checkpoint_path = "/tmp/finetune_alexnet/"

# Create parent path if it doesn't exist
if not os.path.isdir(checkpoint_path): os.mkdir(checkpoint_path)

我随便选择了最后两层(fc7和fc8)来进行微调。大家可以根据你们的数据集大小选择最后网络层的数量。我的选择可能并不好，但是我只是在这展示一下怎么去选择多个网络。dropout概率的数值和学习率大小还是按照原来网络中的大小设定，你可以视情况进行更改。好好设置这些参数和数据集来得到最好的结果。

现在到一些TensorFlow的东西了。在训练开始之前，我们在TensorFlow里需要一些东西。首先我们需要一些用于输入和标签的占位变量，还有dropout率(in test mode we deactivate dropout, while TensorFlow takes care of activation scaling)。

x = tf.placeholder(tf.float32, [batch_size, 227, 227, 3])
y = tf.placeholder(tf.float32, [None, num_classes])
keep_prob = tf.placeholder(tf.float32)

有了这些，我们可以创建AlexNet对象并定义一个指向模型分数的变量（网络的最后一层，即fc8层）。

# Initialize model
model = AlexNet(x, keep_prob, num_classes, train_layers)

#link variable to model output
score = model.fc8

然后是训练网络需要的计算点。

# List of trainable variables of the layers we want to train
var_list = [v for v in tf.trainable_variables() if v.name.split('/')[0] in train_layers]

# Op for calculating the loss
with tf.name_scope("cross_ent"):
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits = score, labels = y))

# Train op
with tf.name_scope("train"):
    # Get gradients of all trainable variables
    gradients = tf.gradients(loss, var_list)
    gradients = list(zip(gradients, var_list))

    # Create optimizer and apply gradient descent to the trainable variables
    optimizer = tf.train.GradientDescentOptimizer(learning_rate)
    train_op = optimizer.apply_gradients(grads_and_vars=gradients)

# Add gradients to summary
for gradient, var in gradients:
    tf.summary.histogram(var.name + '/gradient', gradient)

# Add the variables we train to the summary
for var in var_list:
    tf.summary.histogram(var.name, var)

# Add the loss to summary
tf.summary.scalar('cross_entropy', loss)

相比于在Keras里训练模型所要做的工作（在模型上调用fit()），这些东西乍看起来非常复杂，很难的样子，但是注意，这是TensorFlow。Train op本可以简单一些（用optimizer.minimize()），但是像这样的话，要是想要知道梯度是否经过了我们想要训练的网络层，我们就能获取梯度并让它们在TensorBoard中展示出来，这不是很酷么！接下来我们定义一个评估的计算点——精度。

# Evaluation op: Accuracy of the model
with tf.name_scope("accuracy"):
    correct_pred = tf.equal(tf.argmax(score, 1), tf.argmax(y, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

# Add the accuracy to the summary
tf.summary.scalar('accuracy', accuracy)

训练前最后的一点工作就是初始化tf.FileWriter和tf.train.Saver用于模型持久化，还有初始化图片生成器对象。

# Merge all summaries together
merged_summary = tf.summary.merge_all()

# Initialize the FileWriter
writer = tf.summary.FileWriter(filewriter_path)

# Initialize an saver for store model checkpoints
saver = tf.train.Saver()

# Initalize the data generator seperately for the training and validation set
train_generator = ImageDataGenerator(train_file, horizontal_flip = True, shuffle = True)
val_generator = ImageDataGenerator(val_file, shuffle = False)

# Get the number of training/validation steps per epoch
train_batches_per_epoch = np.floor(train_generator.data_size / batch_size).astype(np.int16)
val_batches_per_epoch = np.floor(val_generator.data_size / batch_size).astype(np.int16)

好了，现在是循环训练：想想该怎么做？我们会建立一个TensorFlow会话，初始化所有变量，加载所有不想进行训练网络层的预训练参数，然后进行循环训练操作。每一步我们会用FileWriter存储一些变量值概览。每个epoch之后我们会评估模型并保存模型检查点。这里运行如下代码

# Start Tensorflow session
with tf.Session() as sess:

    # Initialize all variables
    sess.run(tf.global_variables_initializer())

    # Add the model graph to TensorBoard
    writer.add_graph(sess.graph)

    # Load the pretrained weights into the non-trainable layer
    model.load_initial_weights(sess)

    print("{} Start training...".format(datetime.now()))
    print("{} Open Tensorboard at --logdir {}".format(datetime.now(), filewriter_path))

    # Loop over number of epochs
    for epoch in range(num_epochs):

        print("{} Epoch number: {}".format(datetime.now(), epoch+1))

        step = 1

        while step < train_batches_per_epoch:

            # Get a batch of images and labels
            batch_xs, batch_ys = train_generator.next_batch(batch_size)

            # And run the training op
            sess.run(train_op, feed_dict={x: batch_xs, y: batch_ys, keep_prob: dropout_rate})

            # Generate summary with the current batch of data and write to file
            if step%display_step == 0:
                s = sess.run(merged_summary, feed_dict={x: batch_xs, y: batch_ys, keep_prob: 1.})
                writer.add_summary(s, epoch*train_batches_per_epoch + step)

            step += 1

        # Validate the model on the entire validation set
        print("{} Start validation".format(datetime.now()))
        test_acc = 0.
        test_count = 0
        for _ in range(val_batches_per_epoch):
            batch_tx, batch_ty = val_generator.next_batch(batch_size)
            acc = sess.run(accuracy, feed_dict={x: batch_tx, y: batch_ty, keep_prob: 1.})
            test_acc += acc
            test_count += 1
        test_acc /= test_count
        print("Validation Accuracy = {:.4f}".format(datetime.now(), test_acc))

        # Reset the file pointer of the image data generator
        val_generator.reset_pointer()
        train_generator.reset_pointer()

        print("{} Saving checkpoint of model...".format(datetime.now()))

        #save checkpoint of the model
        checkpoint_name = os.path.join(checkpoint_path, 'model_epoch'+str(epoch)+'.ckpt')
        save_path = saver.save(sess, checkpoint_name)

        print("{} Model checkpoint saved at {}".format(datetime.now(), checkpoint_name))

这就完成了。让我们看看训练过程中的准确度和损失函数图吧。我们可以看到，准确度一开始约为50%（正常的），然后在训练集上训练很快达到了95%的准确度。第一个epoch之后验证机准确度为0.9545。

Screenshot of the training process visualized with TensorBoard. At the top is the accuracy, at the bottom the cross-entropy-loss.

模型能够这么快达到一个如此好的准确率是因为我所选择的数据集例子——猫和狗。用于AlexNet训练的ImageNet数据库本身就已经包括了很多不同种类的狗和猫。但是不管怎样，只要改几行代码就能得到一个一般性的AlexNet模型，可以用来解决自己遇到的一些问题。

如果你想要从任何你的检查点开始继续训练，你只要把model.load_initial_weights(sess)这行代码改成下面的代码。

# To continue training from one of your checkpoints
saver.restore(sess, "/path/to/checkpoint/model.ckpt")

最后的一些Notes

在使用脚本的过程中，大家没有必要一定要用我的ImageDataGenerator类（这个类不一定很高效）。大家大可以使用自己的方式来为训练网络模型提供图片和标签数据。

如果你要任何问题，可以随意提问。再说一遍，所有代码都可以在github上找到。但是要注意，就像开头说的，由于TensorFlow 1.12rc0的新的输入方式，我更新了所有代码。如果你想要用更新后的版本，请先确认一下你的TensorFlow版本是否更新了。