在 CIFAR10 小型图像数据集上训练一个简单的 CNN-Capsule Network。

无数据增益的情况下：在 10 轮迭代后验证集准确率达到 75%，在 15 轮后达到 79%，在 20 轮后过拟合。

有数据增益情况下:在 10 轮迭代后验证集准确率达到 75%，在 15 轮后达到 79%，在 30 轮后达到 83%。在我的测试中，50 轮后最高的验证集准确率为 83.79%。

这是一个快速版的实现，在 GTX 1070 GPU 上迭代只需 20s/epoch。

  1. from __future__ import print_function
  2. from keras import backend as K
  3. from keras.layers import Layer
  4. from keras import activations
  5. from keras import utils
  6. from keras.datasets import cifar10
  7. from keras.models import Model
  8. from keras.layers import *
  9. from keras.preprocessing.image import ImageDataGenerator
  12. # 挤压函数
  13. # 我们在此使用 0.5，而不是 Hinton 论文中给出的 1
  14. # 如果为 1，则向量的范数将被缩小。
  15. # 如果为 0.5，则当原始范数小于 0.5 时，范数将被放大，
  16. # 当原始范数大于 0.5 时，范数将被缩小。
  17. def squash(x, axis=-1):
  18.     s_squared_norm = K.sum(K.square(x), axis, keepdims=True) + K.epsilon()
  19.     scale = K.sqrt(s_squared_norm) / (0.5 + s_squared_norm)
  20.     return scale * x
  23. # 使用自定义的 softmax 函数，而非 K.softmax，
  24. # 因为 K.softmax 不能指定轴。
  25. def softmax(x, axis=-1):
  26.     ex = K.exp(x - K.max(x, axis=axis, keepdims=True))
  27.     return ex / K.sum(ex, axis=axis, keepdims=True)
  30. # 定义 margin loss，类似于 hinge loss
  31. def margin_loss(y_true, y_pred):
  32.     lamb, margin = 0.5, 0.1
  33.     return K.sum(y_true * K.square(K.relu(1 - margin - y_pred)) + lamb * (
  34.         1 - y_true) * K.square(K.relu(y_pred - margin)), axis=-1)
  37. class Capsule(Layer):
  38.     """一个由纯 Keras 实现的 Capsule 网络。
  39.     总共有两个版本的 Capsule。
  40.     一个类似于全连接层 (用于固定尺寸的输入)，
  41.     另一个类似于时序分布全连接层 (用于变成输入)。
  43.     Capsure 的输入尺寸必须为 (batch_size,
  44.                              input_num_capsule,
  45.                              input_dim_capsule
  46.                             )
  47.     以及输出尺寸必须为 (batch_size,
  48.                       num_capsule,
  49.                       dim_capsule
  50.                      )
  52.     Capsule 实现来自于 https://github.com/bojone/Capsule/
  53.     Capsule 论文: https://arxiv.org/abs/1710.09829
  54.     """
  56.     def __init__(self,
  57.                  num_capsule,
  58.                  dim_capsule,
  59.                  routings=3,
  60.                  share_weights=True,
  61.                  activation='squash',
  62.                  **kwargs):
  63.         super(Capsule, self).__init__(**kwargs)
  64.         self.num_capsule = num_capsule
  65.         self.dim_capsule = dim_capsule
  66.         self.routings = routings
  67.         self.share_weights = share_weights
  68.         if activation == 'squash':
  69.             self.activation = squash
  70.         else:
  71.             self.activation = activations.get(activation)
  73.     def build(self, input_shape):
  74.         input_dim_capsule = input_shape[-1]
  75.         if self.share_weights:
  76.             self.kernel = self.add_weight(
  77.                 name='capsule_kernel',
  78.                 shape=(1, input_dim_capsule,
  79.                        self.num_capsule * self.dim_capsule),
  80.                 initializer='glorot_uniform',
  81.                 trainable=True)
  82.         else:
  83.             input_num_capsule = input_shape[-2]
  84.             self.kernel = self.add_weight(
  85.                 name='capsule_kernel',
  86.                 shape=(input_num_capsule, input_dim_capsule,
  87.                        self.num_capsule * self.dim_capsule),
  88.                 initializer='glorot_uniform',
  89.                 trainable=True)
  91.     def call(self, inputs):
  92.         """遵循 Hinton 论文中的路由算法，
  93.         但是将 b = b + <u,v> 替换为 b = <u,v>。
  95.         这一改变将提升 Capsule 的特征表示能力。
  97.         然而，你仍可以将
  98.             b = K.batch_dot(outputs, hat_inputs, [2, 3])
  99.         替换为
  100.             b += K.batch_dot(outputs, hat_inputs, [2, 3])
  101.         来实现一个标准的路由。
  102.         """
  104.         if self.share_weights:
  105.             hat_inputs = K.conv1d(inputs, self.kernel)
  106.         else:
  107.             hat_inputs = K.local_conv1d(inputs, self.kernel, [1], [1])
  109.         batch_size = K.shape(inputs)[0]
  110.         input_num_capsule = K.shape(inputs)[1]
  111.         hat_inputs = K.reshape(hat_inputs,
  112.                                (batch_size, input_num_capsule,
  113.                                 self.num_capsule, self.dim_capsule))
  114.         hat_inputs = K.permute_dimensions(hat_inputs, (0, 2, 1, 3))
  116.         b = K.zeros_like(hat_inputs[:, :, :, 0])
  117.         for i in range(self.routings):
  118.             c = softmax(b, 1)
  119.             o = self.activation(K.batch_dot(c, hat_inputs, [2, 2]))
  120.             if i < self.routings - 1:
  121.                 b = K.batch_dot(o, hat_inputs, [2, 3])
  122.                 if K.backend() == 'theano':
  123.                     o = K.sum(o, axis=1)
  125.         return o
  127.     def compute_output_shape(self, input_shape):
  128.         return (None, self.num_capsule, self.dim_capsule)
  131. batch_size = 128
  132. num_classes = 10
  133. epochs = 100
  134. (x_train, y_train), (x_test, y_test) = cifar10.load_data()
  136. x_train = x_train.astype('float32')
  137. x_test = x_test.astype('float32')
  138. x_train /= 255
  139. x_test /= 255
  140. y_train = utils.to_categorical(y_train, num_classes)
  141. y_test = utils.to_categorical(y_test, num_classes)
  143. # 一个常规的 Conv2D 模型
  144. input_image = Input(shape=(None, None, 3))
  145. x = Conv2D(64, (3, 3), activation='relu')(input_image)
  146. x = Conv2D(64, (3, 3), activation='relu')(x)
  147. x = AveragePooling2D((2, 2))(x)
  148. x = Conv2D(128, (3, 3), activation='relu')(x)
  149. x = Conv2D(128, (3, 3), activation='relu')(x)
  152. """现在我们将其尺寸重新调整为 (batch_size, input_num_capsule, input_dim_capsule)，再连接一个 Capsule 网络。
  154. 最终模型的输出为长度为 10 的 Capsure，其 dim=16。
  156. Capsule 的长度表示为 proba，
  157. 因此问题变成了一个『10个二分类』的问题。
  158. """
  160. x = Reshape((-1, 128))(x)
  161. capsule = Capsule(10, 16, 3, True)(x)
  162. output = Lambda(lambda z: K.sqrt(K.sum(K.square(z), 2)))(capsule)
  163. model = Model(inputs=input_image, outputs=output)
  165. # 使用 margin loss
  166. model.compile(loss=margin_loss, optimizer='adam', metrics=['accuracy'])
  167. model.summary()
  169. # 可以比较有无数据增益对应的性能
  170. data_augmentation = True
  172. if not data_augmentation:
  173.     print('Not using data augmentation.')
  174.     model.fit(
  175.         x_train,
  176.         y_train,
  177.         batch_size=batch_size,
  178.         epochs=epochs,
  179.         validation_data=(x_test, y_test),
  180.         shuffle=True)
  181. else:
  182.     print('Using real-time data augmentation.')
  183.     # 这一步将进行数据处理和实时数据增益:
  184.     datagen = ImageDataGenerator(
  185.         featurewise_center=False,  # 将整个数据集的均值设为 0
  186.         samplewise_center=False,  # 将每个样本的均值设为 0
  187.         featurewise_std_normalization=False,  # 将输入除以整个数据集的标准差
  188.         samplewise_std_normalization=False,  # 将输入除以其标准差
  189.         zca_whitening=False,  # 运用 ZCA 白化
  190.         zca_epsilon=1e-06,  # ZCA 白化的 epsilon值
  191.         rotation_range=0,  # 随机旋转图像范围 (角度, 0 to 180)
  192.         width_shift_range=0.1,  # 随机水平移动图像 (总宽度的百分比)
  193.         height_shift_range=0.1,  # 随机垂直移动图像 (总高度的百分比)
  194.         shear_range=0.,  # 设置随机裁剪范围
  195.         zoom_range=0.,  # 设置随机放大范围
  196.         channel_shift_range=0.,  # 设置随机通道切换的范围
  197.         # 设置填充输入边界之外的点的模式
  198.         fill_mode='nearest',
  199.         cval=0.,  # 在 fill_mode = "constant" 时使用的值
  200.         horizontal_flip=True,  # 随机水平翻转图像
  201.         vertical_flip=False,  # 随机垂直翻转图像
  202.         # 设置缩放因子 (在其他转换之前使用)
  203.         rescale=None,
  204.         # 设置将应用于每一个输入的函数
  205.         preprocessing_function=None,
  206.         # 图像数据格式，"channels_first" 或 "channels_last" 之一
  207.         data_format=None,
  208.         # 保留用于验证的图像比例（严格在0和1之间）
  209.         validation_split=0.0)
  211.     # 计算特征标准化所需的计算量
  212.     # (如果应用 ZCA 白化，则为 std，mean和主成分)。
  213.     datagen.fit(x_train)
  215.     # 利用由 datagen.flow() 生成的批来训练模型。
  216.     model.fit_generator(
  217.         datagen.flow(x_train, y_train, batch_size=batch_size),
  218.         epochs=epochs,
  219.         validation_data=(x_test, y_test),
  220.         workers=4)