前言

好久不见！这几天虽然放假在家，但是各种事情还是非常多，忙里偷闲做了一个给动漫图片添加tag的神经网络模型。因此写下这篇文章来记录一下过程。

给动漫图片添加tag属于多标签分类问题：一张图片可以有多个标签。

这类问题不能简单使用准确率来判断模型的性能，因此我用通过召回率和精确率进行计算的F1值来评价模型的性能。

最终结果：模型在 danbooru网站上的20万张图片上进行训练，能够在470个tag中选择图片合适的tag，F1值达到0.8。

在训练时我们使用了 Google Colab平台的TPU，选择这个方案的原因是他的速度确实比Colab上的GPU快很多。

项目详细

网络设计

根据这类问题的特性：一张图片可能会有多个标签，因此不可以使用 softmax激活函数来作为模型输出层的激活函数。如果要列举每一种可能的标签组合来使用 softmax的话，那输出层的大小是不可以接受的。因此使用 sigmoid激活函数

我们的网络基于 ResNet50，但对输出层进行了一些更改。下面是网络结构

1
2
3
4
5
6
7
8
base_model = keras.applications.resnet.ResNet50(include_top=False, weights=None, input_shape=(SIZE, SIZE, 3))
model = keras.Sequential([
      base_model,
      keras.layers.Conv2D(filters=len(tags),kernel_size=(1,1),padding='same'),
      keras.layers.BatchNormalization(epsilon=1.001e-5),
      keras.layers.GlobalAveragePooling2D(name='avg_pool'),
      keras.layers.Activation("sigmoid")
    ])

其中输入大小是 224*224*3

之所以采用这样的结构是因为在输出结果的卷积层后面加上标准化可以加快模型的收敛速度

数据预处理

以下是数据预处理的代码，我的数据预处理方案是先对图片使用各种随机变换，再将图片等比缩放到SIZE大小，空白位置填充黑色到 SIZE*SIZE大小

1
2
3
4
5
6
7
8
9
10
11
12
13
14
@tf.function()
def _load_train(x):
    img = tf.io.decode_jpeg(x["c"], channels=3)
    img = tf.image.random_flip_left_right(img)
    img = tf.image.random_flip_up_down(img)
    img = tf.image.random_contrast(img, 0.8, 1.2)
    img = tf.image.random_brightness(img, 0.2)
    img = tf.image.random_saturation(img, 0.8, 1.2)
    img = tf.image.random_hue(img, 0.05)
    img = tf.image.resize_with_pad(img, SIZE, SIZE, method="nearest")
    img = img / 255
    tag = x["t"]
    return img , tag

这里需要注意 resize_with_pad的 method非常关键。如果使用默认的方法可能会导致图片出现错误，大幅影响模型的性能。

下面是预处理的输出图片，同一张图片可以生成很多张不同的图片

img

训练

如果使用 BinaryCrossentropy 作为损失函数会导致模型收敛速度过慢，这是数据集内各类别的样本不均衡导致的，因此我使用 SigmoidFocalCrossEntropy 来作为训练的损失函数，关于该函数的论文 https://arxiv.org/abs/1708.02002

这个函数的特性在于可以使已经正确分类的类别贡献的损失值减少，使模型的训练过程更专注于训练未能正确分类的类别

为了达到更好的模型性能，我使用了 SGD优化器，并使用 cosine 函数来控制学习率，如图所示

batch_size设置为256来用完所有TPU内存

实现详细

其实现在才开始这篇文章的重点：如何充分使用Colab上的免费TPU？

首先Colab并不能一直运行你的代码，它时不时会给你弹出验证码，但是在一两次之后就不会再弹出验证码了，你可以在此之后随便使用。当然在你的使用量过多之后可能会无法分配TPU实例。

其次，在TPU上训练模型时你的数据只能从RAM或者是 Google Cloud Storage上，并需要将你的数据打包成 TFRecord来快速读取

当然，你的代码也需要一些调整。下面分块来介绍

在TPU上启动训练

可以参考官方的示例来快速开始 官方链接

关键点就只有下面几步

  1. 连接和初始化

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    import tensorflow as tf
    print("Tensorflow version " + tf.__version__)

    try:
      tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection
      print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
    except ValueError:
      raise BaseException('ERROR: Not connected to a TPU runtime; please see the previous cell in this notebook for instructions!')

    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    tpu_strategy = tf.distribute.experimental.TPUStrategy(tpu)

  2. 在TPU上创建模型

    1
    2
    with tpu_strategy.scope(): # creating the model in the TPUStrategy scope means we will train the model on the TPU
      model = create_model()

  3. 开始训练

    1
    2
    history = model.fit(training_dataset, validation_data=validation_dataset,
                        steps_per_epoch=train_steps, epochs=EPOCHS, callbacks=[lr_callback])

是不是看起来很简单？但是这只是其中最基础的使用方式，如果需要自定义训练还需要添加更多代码

可以看到官方示例就是使用 TFRecord 格式的，下面快速介绍 TFRecord 格式的打包和读取

TFRecord的打包和读取

关于该格式的官方文档 官方链接

打包代码如下

1
2
3
4
5
6
writer = tf.io.TFRecordWriter("xxx.tfrd", TFRecordCompressionType.ZLIB)
feature = {
                "c": tf.train.Feature(bytes_list=tf.train.BytesList(value=[content])),
                "t": tf.train.Feature(int64_list=tf.train.Int64List(value=tag))
           }
writer.write(tf.train.Example(features=tf.train.Features(feature=feature)).SerializeToString())

这段代码将图片编码为jpeg后的二进制和编码后的tag存储到xxx.tfrd内，并使用ZLIB压缩

其中 tag 是python列表，列表内存放整数，形如 [0,0,0,1] ，注意这里必须使用python的数据类型，不能使用numpy或者是tf的

其中 c 和 t 可以自己选择名称，写入时不能中断，否则重新启动就会从头开始写入

加载代码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

dataset = tf.data.TFRecordDataset("xxx.tfrd", compression_type="ZLIB")

image_feature_description = {
    'c': tf.io.FixedLenFeature([], tf.string),
    't': tf.io.FixedLenFeature([length], tf.int64),
}

@tf.function()
def _load(x):
    img = tf.io.decode_jpeg(x["c"], channels=3)
    img = img / 255
    return img

@tf.function
def _parse(proto):
    return tf.io.parse_single_example(proto, image_feature_description)

dataset2 = dataset.map(_parse).map(_load)

这样就能将数据集加载出来了，其中 c 和 t 和上面设置的一样。length对应上面的python列表的长度。所有的列表应该是一样长的，否则就需要另一种加载方式

在这之后就能使用上面提到的数据预处理方案进行操作了

Cloud Storage的使用

这里提供一种免费使用 Cloud Storage 的方案， 这种方案可以免费存储 5GB 的数据，并获得少量数据传输配额，足够在colab上学习深度学习了

首先注册 FireBase ，创建一个项目并进入控制台，在左侧选择 Storage

开通 Storage，一路下一步即可

使用相同的账号登录 Google Cloud Console ，选择和刚刚Firebase相同的项目，转到 Cloud Storage；这时候你就会看到就算未绑定卡也有两个存储账户

打开和Firebase相同的那一个，往里面传文件。

传好之后设置权限，在右上角找到 Cloud Shell 。敲下面的命令将所有文件设置成公开可读取

1
gsutil acl ch -g ALL:R 路径**

其中，路径可通过Firebase获取(gs://开头的链接) ，或者是在控制台右键文件获取文件链接之后将文件名改成 **

之后将 gs:// 开头的 TFRecord文件路径填写到上面读取的代码里就可以了，如果有多个文件可以传递列表

在TPU上自定义训练

和本地的自定义训练差不多，只是需要做一些更改，下面是例子。更多内容查阅官方文档即可，这里没什么坑

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
with tpu_strategy.scope():
  loss = keras.losses.BinaryCrossentropy(reduction="none")
  optim = keras.optimizers.Adam()
  meanError = keras.metrics.MeanAbsoluteError()


@tf.function
def distributed_train_step(x,y):
  def train(x,y):
    with tf.GradientTape() as tape:
      out = model(x,training=True)
      l = loss(y,out)
      l = tf.nn.compute_average_loss(l,global_batch_size=128)
    grad = tape.gradient(l,model.trainable_variables)
    optim.apply_gradients(zip(grad,model.trainable_variables))
    return l * tpu_strategy.num_replicas_in_sync
  
  l = tpu_strategy.run(train, args=(x,y,))
  return tpu_strategy.reduce(tf.distribute.ReduceOp.SUM, l,
                         axis=None)

这是一个需要每一个批次训练之后的loss传递会调用方的例子，如果不需要传回可以删除return以及后面的内容。

使用TensorBoard

和本地情况不一样，TPU的TensorBoard也只能存数据到Cloud Storage

将TensorBoard的callback里的路径改成 gs://xxxx之后还需要在colab运行下面的代码

1
!gcloud auth application-default login

并且需要安装TPU支持来查看数据

1
2
pip install tensorboard-plugin-profile
pip install --upgrade "cloud-tpu-profiler"

使用感受

感受就是。。。。。。。确实蛮快的，比学校里面的V100还快一些，180TFLOPS的算力应该不是瞎吹的，在 ResNet50上一秒钟可以处理1000多张 224*224 样本图片

此外试了下朋友的Colab Pro账户，可以毫无中断的使用TPU训练超过12小时以上（前提是你得把浏览器挂电脑上），也不会有验证码的打扰，也没有遇到配额限制