Caffe2 教程-CIFAR 10

发表于   |   更新于 | 分类于 | 评论数： | 阅读次数：
本文字数： 37k | 阅读时长 ≈ 34 分钟

前言

本示例主要包含以下几个主要部分:

  • 下载 Cifar10 数据集
  • 将 Images 写入 lmdbs
  • 定义并且训练 model
  • 保存训练好的 model
  • 加载训练好的 model
  • 在 testing imdb 上面执行 inference
  • 继续训练以提高 test accuracy
  • 测试 retrained model

项目地址: https://github.com/caffe2/tutorials

导入必要的包

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

%matplotlib inline
from matplotlib import pyplot as plt
import numpy as np
import os
import lmdb
import shutil
from imageio import imread
import caffe2.python.predictor.predictor_exporter as pe
from caffe2.proto import caffe2_pb2
from caffe2.python.predictor import mobile_exporter
from caffe2.python import(
    brew,
    core,
    model_helper,
    net_drawer,
    optimizer,
    visualize,
    workspace
)

# 设置全局初始化信息 , 如果需要看到更加详细的初始化信息, 需要将caffe2_log_level参数设置为 1
core.GlobalInit(['caffe2', '--caffe2_log_level=0'])

下载并解压数据集

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import requests
import tarfile
# 设置下载路径相关变量
# data_folder为data下载并解压的地方, 这个路径需要根据你自己的电脑环境设置
data_folder = os.path.join(
            os.path.expanduser('~'),
            'caffe2_notebooks', 'tutorial_data', 'cifar10'
            )
# root_folder为checkpoints文件和.pb模型定义文件存放的地方
root_folder = os.path.join(
            os.path.expanduser('~'),
            'caffe2_notebooks', 'tutorial_files', 'tutorial_cifar10'
            )

url = "http://pjreddie.com/media/files/cifar.tgz"   # url to data
filename = url.split("/")[-1]                       # download file name
download_path = os.path.join(data_folder, filename) # path to extract data to

if not os.path.isdir(data_folder):
    os.makedirs(data_folder)

# 如果data不存在, 就下载并解压之
if not os.path.exists(download_path.strip('.tgz')):
    # Download data
    r = requests.get(url, stream=True)
    print("Downloading... {} to {}".format(url, download_path))
    open(download_path, 'wb').write(r.content)
    print("Finished downloading...")

    # Unpack images from tgz file
    print('Extracting images from tarball...')
    tar = tarfile.open(download_path, 'r')
    for item in tar:
        tar.extract(item, data_folder)
    print("Completed download and extraction!")

else:
    print("Image directory already exists. Moving on...")

接下来看一看数据集中的一些示例:

1
2
3
4
5
6
7
8
9
10
11
12
13
import glob
from skimage import io

# Grab 5 image paths from training set to display
sample_imgs = glob.glob(os.path.join(data_folder, "cifar", "train") + '/*.png')[:5]

# Plot images
f, ax = plt.subplots(1, 5, figsize=(10,10))
plt.tight_layout()
for i in range(5):
    ax[i].set_title(sample_imgs[i].split("_")[-1].split(".")[0])
    ax[i].axis('off')
    ax[i].imshow(io.imread(sample_imgs[i]).astype(np.uint8))

创建 LMDBs 数据结构

现在我们已经拥有了数据, 接下来需要写入LMDBs用于训练, 验证和测试. 为了根据每个类别来分离图片, 下面将会使用到一个经常在caffe框架中使用的技术, 创建label files. 具体来说, label files就是将每个.png图片对应到它的类别上面去:
/path/to/im1.png 7
/path/to/im2.png 3
/path/to/im3.png 5
/path/to/im4.png 0
…

根据不同的数据集, 创建labels的方式有些许区别

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
# Paths to train and test directories
training_dir_path = os.path.join(
            os.path.expanduser('~'),
            'caffe2_notebooks', 'tutorial_data', 'cifar10', 'cifar', 'train'
            )
testing_dir_path = os.path.join(
            os.path.expanduser('~'),
            'caffe2_notebooks', 'tutorial_data', 'cifar10', 'cifar', 'test'
            )

# Paths to label files
training_labels_path = os.path.join(
            os.path.expanduser('~'),
             'caffe2_notebooks', 'tutorial_data', 'cifar10', 'training_dictionary.txt'
              )
validation_labels_path = os.path.join(
            os.path.expanduser('~'),
            'caffe2_notebooks', 'tutorial_data', 'cifar10', 'validation_dictionary.txt'
            )
testing_labels_path = os.path.join(
            os.path.expanduser('~'),
            'caffe2_notebooks', 'tutorial_data', 'cifar10', 'testing_dictionary.txt'
            )

# Paths to LMDBs
training_lmdb_path = os.path.join(
            os.path.expanduser('~'),
            'caffe2_notebooks', 'tutorial_data', 'cifar10', 'training_lmdb'
            )
validation_lmdb_path = os.path.join(
            os.path.expanduser('~'),
            'caffe2_notebooks', 'tutorial_data', 'cifar10', 'validation_lmdb'
            )
testing_lmdb_path = os.path.join(
            os.path.expanduser('~'),
            'caffe2_notebooks', 'tutorial_data', 'cifar10', 'testing_lmdb'
            )

# Path to labels.txt
labels_path = os.path.join(
            os.path.expanduser('~'),
            'caffe2_notebooks', 'tutorial_data', 'cifar10', 'cifar', 'labels.txt'
            )

# Open label file handler
labels_handler = open(labels_path, "r")


# Create classes dictionary to map string labels to integer labels
classes = {}
i = 0
lines = labels_handler.readlines()
for line in sorted(lines):
    line = line.rstrip()
    classes[line] = i
    i += 1
labels_handler.close()

print("classes:", classes)

输出如下:

1
2
3
4
classes: {
    'automobile': 1, 'airplane': 0, 'truck': 9, 'ship': 8, 'dog': 5,
    'horse': 7, 'cat': 3, 'frog': 6, 'deer': 4, 'bird': 2
}

既然现在我们已经拥有了从类别字符串名到数字的映射关系, 我们就可以创建用于训练, 验证和推演的标签文件了. 我们会将数据分成以下三个部分:

  • training: 44000 images (73%)
  • validation: 6000 images (10%)
  • testing: 10000 image (17%)

注意到我们的验证数据集仅仅只是训练数据集的一个子集, 我们这么做是为了查看我们当前的模型在面对没有见过的数据集时的表现怎么样. 为了得到分布相对均匀的训练集合验证集, 我们首先读取所有的图片(路径), 将其随机打乱, 然后拆分成数据集和验证集. 代码如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
from random import shuffle

# Open file handlers
training_labels_handler = open(training_labels_path, "w")
validation_labels_handler = open(validation_labels_path, "w")
testing_labels_handler = open(testing_labels_path, "w")

# 创建 training, validation, testing 的标签文件
i = 0
validation_count = 6000
# 将所有的训练集的图片路径读取到imgs中, 以字符串列表的形式存储
imgs = glob.glob(training_dir_path+"/*.png")
shuffle(imgs) # 打乱文件路径
for img in imgs:
    if i < validation_count:
        # img的形式为 path/train/31205_frog.png
        validation_labels_handler.write(img + ' ' + str(classes[img.split('_')[-1].split('.')[0]]) + '\n')
    else:
        training_labels_handler.write(img + ' ' + str(classes[img.split('_')[0].split('.')[0]]) + '\n')
    i += 1
print("Finished writing training and validation label files")

# 根据 test image 的信息存储 test labels 到文件中
for img in glob.glob(testing_dir_path+'/*.png'): # 获取所有 test 图片的文件路径
    testing_labels_handler.write(img + ' ' + str(classes[img.split('_')[-1].split('.')[0]]) + '\n')
print("Finished writing testing label files")

# Close file handlers
training_labels_handler.close()
validation_labels_handler.close()
testing_labels_handler.close()

现在, 我们已经可以开始用这些创建好的标签文件来构建我们的 LMDBs 数据集了. 下面的代码采纳自 Caffe2 的 lmdb_create_example.py 脚本. 请注意在将数据送入 LMDB 之前, 首先要将图片的颜色通道从 RGB 变成 BGR, 并且要将图片矩阵从 HWC 变成 CHW. 另外, Caffe2 接受的输入格式为 NCHW, 这里的 N 代表 batch-size.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def write_lmdb(labels_file_path, lmdb_path):
    labels_handler = open(labels_file_path, "r")
    print(">>>Write database...")
    LMDB_MAP_SIZE = 1 << 40 # 设置 lmdb 的size
    print("LMDB_MAP_SIZE", LMDB_MAP_SIZE)
    env = lmdb.open(lmdb_path, map_size=LMDB_MAP_SIZE) # <class 'Environment'>

    with env.begin(write=True) as txn:
        count = 0
        for line in labels_handler.readlines():
            line = line.rstrip() # 去除最后的回车
            im_path = line.split()[0]
            im_label = int(line.split()[1])

            # from imageio import imread, 读取图片
            img_data = imread(im_path).astype(np.float32)
            # 将 RGB 转换成 BGR 格式
            img_data = img_data[:, :, (2,1,0)]
            # 将 HWC 转换成 CHW 格式
            img_data = np.transpose(img_data, (2,0,1))

            tensor_protos = caffe2_pb2.TensorProtos() # 创建 TensorProtos
            img_tensor = tensor_protos.protos.add() # 获取 TensorProto
            img_tensor.dims.extend(img_data.shape) # 将img_tensor.dims从[]扩展成[3,32,32]
            img_tensor.data_type = 1 # 将类型置为 float
            flatten_img = img_data.reshape(np.prod(img_data.shape)) # 将 img_data 的维度变成 32×32×3 = 3072 维(单维度)
            img_tensor.float_data.extend(flatten_img) # 将 flatten_img 填入 img_tensor(初始时为空)
            label_tensor = tensor_protos.protos.add() # 获取一个新的 TensorProto
            label_tensor.data_type = 2 # 置为 int 类型
            label_tensor.int32_data.append(im_label) # 将 label(int类型, 表示标签编号) 添加到 label_tensor 中.
            txn.put( # 第一个参数为 key, 第二个为 value
                '{}'.format(count).encode('ascii'),
                tensor_protos.SerializeToString()
            )
            if ((count % 1000 == 0)):
                print("Inserted {} rows".format(count))
            count = count + 1

        print("Inserted {} rows".format(count))
        print("\nLMDB saved at " + lmdb_path + "\n\n")
        labels_path.close()

# 调用上面定义好的 LMDBs 写入函数
if not os.path.exists(training_lmdb_path):
    print("Writing training LMDB")
    write_lmdb(training_labels_path, training_lmdb_path)
else:
    print(training_lmdb_path, "already exists!")
if not os.path.exists(validation_lmdb_path):
    print("Writing validation LMDB")
    write_lmdb(validation_labels_path, validation_lmdb_path)
else:
    print(validation_lmdb_path, "already exists!")
if not os.path.exists(testing_lmdb_path):
    print("Writing testing LMDB")
    write_lmdb(testing_labels_path, testing_lmdb_path)
else:
    print(testing_lmdb_path, "already exists!")

定义 CNN 模型

既然我们已经将数据存储成了 LMDBs 的格式, 那么是时候建立相应的网络模型了! 首先, 设置一些路径变量, 定义数据集的具体参数, 同时定义训练参数.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# caffe2 的 init 网络 和 predict 网络的位置
init_net_out = 'cifar10_init_net.pb'
predict_net_out = 'cifar10_predict_net.pb'

# 数据集相关参数
image_width = 32
image_height = 32
image_channels = 3
num_classes = 10

# 网络训练参数
training_iters = 2000
training_net_batch_size = 100
validation_images = 6000
validation_interval = 100  # 每经过100次训练迭代, 就会进行一次验证迭代.
checkpoint_iters = 1000 # 每经过1000次迭代, 就会输出一次 checkpoint db

创建相应的目录, 然后将该目录设置为工作目录.

1
2
3
4
5
6
# root_folder = /home/zerozone/caffe2_notebooks/tutorial_files/tutorial_cifar10
if not os.path.isdir(root_folder):
    os.makedirs(root_folder) #

# 将根目录设置为工作目录
workspace.ResetWorkspace(root_folder)

AddInput() 数据输入层

下面的任务是利用 helper functions 来模块化我们的代码, 最终使其构成相应的网络模型. 我们将会使用 ModelHelper 类来定义和表示我们的模型, 同时会包含相应的参数信息. 我们将利用 brew 模块来向我们的模型中添加相应的网络层.
需要注意的是, 调用上面的代码并不会使模型执行任何实际上的计算, 相反的, 我们是在搭建一个计算操作的图结构(graph of operators), 这个图最终指明了当数据传入时, 网络应该进行哪些 forward 和 backward 计算.
第一个 helper function 是AddInput, 它向我们的网络添加了输入(数据)层. 注意到, 图片已经存储到了我们的 LMDBs 数据结构中, 因此在将其送入网络模型之前, 还需要一些简单的预处理操作. 第一, 我们先从 LMDB 中读取 raw image data 和 labels, 它们的类型都是 uint8([0,255] pixel values). 接着我们将数据转换成浮点类型, 同时将图片的像素值放缩到 [0,1] 之间, 使其在训练时可以更快收敛. 最后, 我们会调用model.StopGradient(data, data)来组织在 backward 过程中计算关于数据的梯度(因为我们只需要关于参数的梯度). 最后, 是一些关于引号中的名称含义:

  • “data_unint8” 和 “label” 代表着与 DB 输入数据相关的 blobs 的名字
  • 如果名字是 input blob, 则代表当 operator 运行时的 blob_name
  • 如果名字是 output blob, 则代表是某个 operator 创建的 blob 的名字.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def AddInput(model, batch_size, db, db_dtype):
    # 加载数据
    data_uint8, label = brew.db_input(
        model, # ModelHelper 实例
        blobs_out=["data_uint8", "label"],
        batch_size=batch_size,
        db=db,
        db_type=db_type,
    )

    # 将数据转换成 float
    data = model.Cast(data_uint8, "data", to=core.DataType.FLOAT)
    # 将 data 从 [0, 255] 放缩到 [0,1]
    data = model.Scale(data, data, scale=float(1./256))
    # 避免计算数据的梯度
    data = model.StopGradient(data, data)
    return data, label

根据后文的调用语句:

1
2
3
4
5
data, label = AddInput(
                train_model_zz, batch_size=training_net_batch_size, # 100
                db=training_lmdb_path, # /home/zerozone/caffe2_notebooks/tutorial_data/cifar10/training_lmdb
                db_type="lmdb"
)

调用后, 模型中会增加四个新的 op(原本模型中的 op 为 0 个), 如下所示:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
name: "train_net"
op {
  input: "dbreader_/home/zerozone/caffe2_notebooks/tutorial_data/cifar10/training_lmdb"
  output: "data_uint8"
  output: "label"
  name: ""
  type: "TensorProtosDBInput"
  arg {
    name: "batch_size"
    i: 100
  }
}
op {
  input: "data_uint8"
  output: "data"
  name: ""
  type: "Cast"
  arg {
    name: "to"
    i: 1
  }
}
op {
  input: "data"
  output: "data"
  name: ""
  type: "Scale"
  arg {
    name: "scale"
    f: 0.00390625
  }
}
op {
  input: "data"
  output: "data"
  name: ""
  type: "StopGradient"
}

Add_Original_CIFAR10_Model() 模型核心

处理好输入层以后, 接下来实现 CNN 模型的定义. 本教程使用的网络模型拥有三层卷积层和池化层, 并且使用 ReLU 激活函数. 我们将会使用 update_dims 函数来帮助我们根据卷积层和池化层造成的维度的降低程度. 维度会以下列公式为依据发生变化:

尽管使用 update_dims 函数不是必须的, 但我们发现这是一种避免手动计算维度变化的很不错的策略选择, 尤其是在决定全连接层的参数的时候.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
# 辅助计算维度变化的函数
def update_dims(height, width, kernel, stride, pad):
    new_height = ((height - kernel + 2*pad)//stride) + 1
    new_width = ((width - kernel + 2*pad)//stride) + 1
    return new_height, new_width

def Add_Original_CIFAR10_Model(model, data, num_classes, image_height, image_width, image_channels):
    # 第一层卷积层
    conv1 = brew.conv(
                model, data, 'conv1', dim_in=image_channels, dim_out=32,
                kernel=5, stride=1, pad=2
    )
    h,w = update_dims(height=image_height, width=image_width, kernel=5, stride=1, pad=2)
    # 第一层池化层
    pool1 = brew.max_pool(model, conv1, "pool1", kernel=2, stride=2)
    h,w = update_dims(height=h, width=w, kernel=3, stride=2, pad=0)
    # 第一层激活层
    relu1 = brew.relu(model, pool1, "relu1")

    # 第二层
    conv2 = brew.conv(
                model, relu1, "conv2", dim_in=32, dim_out=32,
                kernel=5, stride=1, pad=2
    )
    h,w = update_dims(height=h, width=w, kernel=5, stride=2, pad=2)
    relu2 = brew.relu(model, conv2, "relu2")
    pool2 = brew.average_pool(model, relu2, "pool2", kernel=3, stride=2)
    h,w = update_dims(height=h, width=w, kernel=3, stride=2, pad=0)

    # 第三次
    conv3 = brew.conv(
                model, pool2, "conv3", dim_in=32, dim_out=64,
                kernel=5, stride=1, pad=2
    )
    h,w = update_dims(height=h, width=w, kernel=5, stride=1, pad=2)
    relu3 = brew.relu(model, conv3, "relu3")
    pool3 = brew.average_pool(model, relu3, "pool3", kernel=3, stride=2)
    h,w = update_dims(height=h, width=w, kernel=3, stride=2, pad=0)

    # 全连接层, 通过动态维护w,h变量, 省去了手动计算的麻烦
    fc1 = brew.fc(model, pool3, "fc1", dim_in=64*h*w, dim_out=64)
    fc2 = brew.fc(model, fc1, "fc2", dim_in=64, dim_out=num_classes)

    # Softmax 网络层, 将输出转换成概率
    softmax = brew.softmax(model, fc2, "softmax")
    return softmax

后文在调用该函数时的代码如下所示:

1
2
3
4
softmax = Add_Original_CIFAR10_Model(
                train_model, data, num_classes
                image_height, image_width, image_channels
)

由于此处涉及的 op 过多, 并且大部分都是重复的, 因此这里我们仅仅贴出网络中 conv1, pool1, relu1 的 op, 其他网络层的 op 类似.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
op {
  input: "data"
  input: "conv1_w"
  input: "conv1_b"
  output: "conv1"
  name: ""
  type: "Conv"
  arg {
    name: "exhaustive_search"
    i: 0
  }
  arg {
    name: "stride"
    i: 1
  }
  arg {
    name: "order"
    s: "NCHW"
  }
  arg {
    name: "pad"
    i: 2
  }
  arg {
    name: "kernel"
    i: 5
  }
  engine: "CUDNN"
}
op {
  input: "conv1"
  output: "pool1"
  name: ""
  type: "MaxPool"
  arg {
    name: "kernel"
    i: 3
  }
  arg {
    name: "cudnn_exhaustive_search"
    i: 0
  }
  arg {
    name: "stride"
    i: 2
  }
  arg {
    name: "order"
    s: "NCHW"
  }
  engine: "CUDNN"
}
op {
  input: "pool1"
  output: "relu1"
  name: ""
  type: "Relu"
  arg {
    name: "cudnn_exhaustive_search"
    i: 0
  }
  arg {
    name: "order"
    s: "NCHW"
  }
  engine: "CUDNN"
}

AddTrainingOperators() 损失函数及优化器

我们的下一个辅助函数是 AddTrainingOperators(). 这个函数将会被我们的训练模型所调用, 用于添加相应的损失函数和优化机制. 我们将会在模型的 softmax 输出和图片的真实标签上使用平均交叉熵来计算模型的损失. 然后我们会添加一个计算损失函数梯度的 operators 到模型中. 最终, 我们会使用 Caffe2 中 optimizer 类的 build_sgd 函数作为我们的优化函数.

1
2
3
4
5
6
7
8
9
10
11
12
def AddTrainingOperators(model, softmax, label):
    xent = model.LabelCrossEntropy([softmax, label], 'xent')
    # 计算 loss 的期望值
    loss = model.AveragedLoss(xent, "loss")
    # from caffe2.python import optimizer
    optimizer.build_sgd(
        model,
        base_learning_rate=0.01,
        policy="fixed",
        momentum=0.9,
        weight_decay=0.004
    )

后文的调用语句如下:

1
AddTrainingOperators(train_model, softmax, label)

此处涉及到的 op 也非常的多, 主要包含了计算之前每一个操作的梯度, 这里我们也只简单的列出一些 op 作为示例:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
op {
  input: "fc1"
  input: "fc2_w"
  input: "fc2_grad"
  output: "fc2_w_grad"
  output: "fc2_b_grad"
  output: "fc1_grad"
  name: ""
  type: "FCGradient"
  arg {
    name: "use_cudnn"
    i: 1
  }
  arg {
    name: "cudnn_exhaustive_search"
    i: 0
  }
  arg {
    name: "order"
    s: "NCHW"
  }
  is_gradient_op: true
}

AddAccuracy() 模型准确率

下面的函数用于计算我们当前模型的 top-1 准确率.

1
2
3
4
5
def AddAccuracy(model, softmax, label):
    # blobs_in 为 [softmax, label]
    # blobs_out 为 "accuracy"
    accuracy = brew.accuracy(model, [softmax, label], "accuracy")
    return accuracy

后文的调用语句如下:

1
AddAccuracy(val_model, softmax, label)

本函数增加了一个新的 op, 如下所示:

1
2
3
4
5
6
7
op {
  input: "softmax"
  input: "label"
  output: "accuracy"
  name: ""
  type: "Accuracy"
}

检查点

下面的函数会在每经过一定次数的迭代之后输出一个 checkpoint db. 一个 checkpoint 本质上来说是模型在训练过程中的一种保存成文件的模型状态. checkpoint 在快速加载训练好或者训练中的模型时非常有用, 并且它们能够在长时间的模型训练的过程中提高有力的安全保障. Caffe2 的 checkpoint 文件类似于 Caffe 的定期输出的 .caffemodel 文件. 我们利用 brew 和 iter operator 来跟踪迭代次数, 并把它们保存成 LMDBs 的形式.
在使用 checkpoints 的时候, 你必须时刻注意在训练过程是否将之前同名检查点给覆盖掉. 如果你尝试覆盖一个 checkpoint db, 那么将会产生报错. 为了解决这种情况, 我们将会把 checkpoints 以一种唯一的命名字典保存在我们的 root_folder 下面. 这个字典的名字会基于当前的系统时间戳, 以避免出现重复.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import datetime

# 在 root_folder 目录下创建唯一的命名字典
unique_timestamp = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))
checkpoint_dir = os.path.join(root_folder. unique_timestamp)
os.makedirs(checkpoint_dir)
print("Checkpoint output location: ", checkpoint_dir)

# 将 checkpoint 添加到给定的模型当中
def AddCheckpoints(train_model, checkpoint_iters, db_type): # 注意这里第一个参数是train_model, 原始教程里写错了, 会导致运行出错
    ITER = brew.iter(train_model, "iter")
    train_model.Checkpoint(
                [ITER] + train_model.params, [],
                db=os.path.join(unique_timestamp, "cifar10_checkpoint_%05d.lmdb"),
                dy_type="lmdb", every=checkpoint_iters
    )

利用 ModelHelper 初始化模型

现在既然我们已经创建了必要的辅助函数, 那么是时候初始化模型, 同时使用相应的函数来定义模型的计算图(operator graphs). 请记住此时我们还没有执行模型.

首先, 定义 train model:

  1. 用 ModelHelper class 初始化模型
  2. 利用 AddInput() 函数添加数据输入层
  3. 添加 Cifar10 模型, 该模型返回一个 softmax blob
  4. 利用 AddTrainingOperators() 函数添加 training operators, 此处需要使用第三步中的 softmax blob
  5. 利用 AddCheckpoints() 函数添加定时的 checkpoints.

下面, 我们来定义 validation model, 该模型在结构上与 training model 是相同的, 但是其数据来源于另一个 LMDB 文件, 并且使用了不同的 batch size. 具体定义如下:

  1. 利用 ModelHelper class 对模型进行初始化(需要将参数 init_params 设置为 False);
  2. 利用 AddInput() 函数添加数据输入层;
  3. 添加 Cifar10 模型, 该模型返回一个 softmax blob;
  4. 利用 AddAccuracy() 函数添加 accuracy layer, 此处需要使用第三步中的 softmax blob;

最后, 我们定义 deploy model 如下:

  1. 利用 ModelHelper class 初始化模型(需令 init_params=False);
  2. 添加 Cifar10 模型, 以 “data” 作为模型输入 input blob.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
arg_scope = {"order": "NCHW"} # 字典, 在 ModelHelper 中, arg_scope["order"] 的默认值即为 "NCHW", 因此该语句可省略

# TRAINING MODEL
# 利用 ModelHelper class 初始化 training model
train_model = model_helper.ModelHelper(name="train_net", arg_scope=arg_scope)
# 从 lmdb 文件添加数据输入层
data, label = AddInput(
    train_model, batch_size=training_net_batch_szie,
    db=training_lmdb_path,
    db_type="lmdb"
)
# 添加模型核心定义, 将返回的值保存在 'softmax' 变量中
softmax = Add_Original_CIFAR10_Model(
                train_model, data, num_classes,
                image_height, image_width, image_channels
)
# 利用上面返回的 softmax 添加 training operators
AddTrainingOperators(train_model, softmax, label)
# 添加 periodic checkpoint 输出到模型中
AddCheckpoints(train_model, checkpoint_iters, db_type="lmdb")

# VALIDATION MODEL
# 利用 ModelHelper class 初始化 val model
val_model = model_helper.ModelHelper(
    name="val_net", arg_scope=arg_scope, init_params=False
)
# 从 validation_lmdb 文件添加数据输入层
data, label = AddInput(
    val_model, batch_size=validation_images,
    db=validation_lmbd_path,
    db_type="lmdb"
)
# 添加模型定义, 并将返回值存储在 'softmax' 变量中
softmax = Add_Original_CIFAR10_Model(
    val_model, data, num_classes,
    image_height, image_width, image_channels
)
# 添加 accuracy operator
AddAccuracy(val_model, softmax, label)

# DEPLOY MODEL
# 利用 ModelHelper class 初始化 deploy model
deploy_model = model_helper.ModelHelper(
    name="deploy_net", arg_scope=arg_scope, init_params=False
)
# 添加模型定义, 输入数据为 "data"
Add_Original_CIFAR10_Model(
    deploy_model, "data", num_classes,
    image_height, image_width, image_channels
)
print("Training, Validation, and Deploy models all defined")

开始训练

最后, 既然我们已经拥有了模型, 并且定义了相应的 operator graphs, 那么就应该开始真正意义上的执行 training process 了. Under the hood, 我们已经定义将我们的模型以 operator graphs 的形式定义, 并且序列化成了 protobuf 格式. 最后一步就是要将这些 protobufs 送入 Caffe2 的 C++ 后端, 以便建立模型并执行它. 还记得 ModelHelper 模型对象拥有的两个网络吗:

  • param_init_net: 包含参数和初始数据
  • net: 包含定义好的网络结构(operator graph)

以上的两个网络都需要执行, 并且我们必须先执行 param_init_net 网络, 注意, 该网络只需要被执行一次, 因此我们利用 workspace.RunNetOnce() 函数来执行, 这个函数会经过一个实例化, 运行, 然后销毁网络的过程. 如果我们需要多次运行某个网络, 就像我们需要多次运行 training nets 和 validation nets 一样, 那么我们必须先利用 workspace.CreateNet() 手动创建网络, 然后利用多次调用 workspace.RunNet() 的方式达到多次运行网络的目的, 从这里我们也就可以看出 RunNetOnce 和 RunNet 之间的区别了, 简单来说, 前者会在执行之前自动创建网络, 而后者只能运行已经创建好的网络, 也正是因为这样, 后者的执行速度要比前者快.
当我们利用 workspace.RunNet() 执行 train_model 以后, 这会在一个 batch 的数据上执行 forward 和 backward 过程, 而在执行 val_model 时, 只会执行 forward 过程(同时会计算准确率).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import math

# 初始化并创建 training network
workspace.RunNetOnce(train_model.param_init_net)
workspace.CreateNet(train_model.net, overwrite=True) # overwrite 参数代表是否允许覆盖checkpoint文件

# 初始化并创建 validation network
workspace.RunNetOnce(val_model.param_init_net)
workspace.CreateNet(val_model.net, overwrite=True)

# 创建 Placeholder, 用于跟踪 loss 和 validation accuracy
loss = np.zeros(int(math.ceil(training_iters/validation_interval))) # shape = (2000/100,) = (20,)
val_accuracy = np.zeros(int(math.ceil(training_iters/validation_interval))) # shape = (20,)
val_count = 0
iteration_list = np.zeros(int(math.ceil(training_iters/validation_interval))) # shape = (20,)

# 现在, 开始运行网络 (forward & backward padd)
for i in range(training_iters):
    workspace.RunNet(train_model.net)

    # 每经过 validation_interval 次迭代就执行一次 validation net
    if (i % validation_interval == 0):
        print("Training iter:", i)
        # 获取当前 workspace 内的名为 "loss" 的 blob
        loss[val_count] = workspace.FetchBlob("loss") # 记录一次 loss
        # 注意, 在 Caffe2 中, 同一个 workspace 的参数会被所有当前 workspace 内的模型所共享,
        # 也就是说, 虽然看起来 train_net 和 val_net 是分别利用 ModelHelper 定义的,
        # 但是, 当他们在workspace中获取参数的值时(通过名字, 如"loss"), 会获得相同的值,
        # 因此, 这里可以直接运行 val_net 来计算 train_net 训练的准确率.
        workspace.RunNet(val_model.net)
        val_accuracy[val_count] = workspace.FetchBlob("accuracy")
        print("Loss:", str(loss[val_count]))
        print("Validation accuracy:", str(val_accuracy[val_count])+"\n")
        iteration_list[val_count] = i # 记录当前迭代次数
        val_count += 1

下面的代码用于将训练过程中的 accuracy 和 loss 绘制出来, 以方便我们观察和分析.

1
2
3
4
5
plt.title("Training loss vs. Validation Accuracy")
plt.plot(iteration_list, loss, 'b')
plt.plot(iteration_list, val_accuracy, 'r')
plt.xlabel("Training iteration")
plt.legend(('Loss', 'Validation Accuracy'), loc='upper right')

保存训练好的模型

当我们在 workspace 中训练好模型的参数以后, 我们就可以利用 mobile_exporter 来导出部署模型. 在 Caffe2 中, 预训练的模型通常情况下会存储成两个单独的 protobuf(.pb) 文件(inti_net 和 predict_net). 模型也可以被存储为 db 格式, 但是最好存储成 pb 格式, 因为这种格式在社区中更流行. 为了保持一致性, 我们将这些参数保存在相同的 checkpoints 的唯一目录下.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 初始化并创建 deploy 网络
workspace.RunNetOnce(deploy_model.param_init_net)
workspace.CreateNet(deploy_model.net, overwrite=True)

# 利用 mobile_exporter 的 Export 函数来获取 init_net 和 predict_net
init_net, predict_net = mobile_exporter.Export(
    workspace, deploy_model.net, deploy_model.params)
)

# 定义 output files 的路径
full_init_net_out = os.path.join(checkpoint_dir, init_net_out)
full_predict_net_out = os.path.join(checkpoint_dir, predict_net_out)

# 将两个网络写入文件
with open(full_init_net_out, 'wb') as f:
    f.write(init_net.SerializeToString())
with open(full_predict_net_out, 'wb') as f:
    f.write(predict_net.SerializeToString())

print("Model saved as " + full_init_net_out + " and " + full_predict_net_out)

利用训练好的模型进行推演

在保存好模型以后, 我们试着来使用这个模型进行推演, 首先定义一些路径相关变量, 方便我们加载特定的模型.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# Train lmdb, 定义训练集的 lmdb路径
TRAIN_LMDB = os.path.join(os.path.expanduser('~'),"caffe2_notebooks/tutorial_data/cifar10/training_lmdb")
# Test lmdb, 定义验证集的 lmdb 路径
TEST_LMDB = os.path.join(os.path.expanduser('~'),"caffe2_notebooks/tutorial_data/cifar10/testing_lmdb")


# 提起保存好的 protobuf 文件.
part1_runs_path = os.path.join(os.path.expanduser('~'), "caffe2_notebooks", "tutorial_files", "tutorial_cifar10")
# runs 为一个列表, 列表的长度跟路径中 checkpoint 的个数有关, 按日期排序, 日期大的在后面
runs = sorted(glob.glob(part1_runs_path + "/*"))

# Init net
INIT_NET = os.path.join(runs[-1], "cifar10_init_net.pb")
# Predict net
PREDICT_NET = os.path.join(runs[-1], "cifar10_predict_net.pb")


# Make sure they all exist
if (not os.path.exists(TRAIN_LMDB)) or (not os.path.exists(TEST_LMDB)) or (not os.path.exists(INIT_NET)) or (not os.path.exists(PREDICT_NET)):
    print("ERROR: input not found!")
else:
    print("Success, you may continue!")

构建 Testing model

1
2
3
4
5
6
7
8
# 利用 ModelHelper 创建 testing model
arg_scope = {"order": "NCHW"}
test_model = model_helper.ModelHelper(
    name="test_model", arg_scope=arg_scope, init_params=False
)

# 添加 data layer, 注意使用 TEST_LMDB
data, _ = AddInputLayer(test_model, 1, TEST_LMDB, 'lmdb') # batch size 为 1

下面我们使用保存好的模型参数来填充(populate) Model Helper. 想要构建用于 testing 的模型, 我们不需要在 model helper 重新创建 params, 同样也不需要添加 gradient operators, 因为我们仅仅执行 forward pass. 我们只需要利用 predict_net.pb 和 init_net.pb 文件中的参数填充 testing model 的 .net 和 .param_init_net 成员即可. 为此, 我们需要从 pb 文件中创建 caffe2_pb 对象, 然后再利用 caffe2_pb 对象创建 Net 对象, 接着将 net 对象 添加(append) 到 .net 和 .param_init 成员中即可. 请注意, 这里的添加操作非常重要, 如果没有执行 append, 我们将会抹去刚刚添加的输入层.

会想上面的程序, 保存好的模型需要一个名为 “data” 输入, 同时会产生一个名为 “softmax” 的输出. Conveniently(but not accidentally), 输入层函数(AddInputLayer())会从 lmdb 中读取数据, 并且将读取到的信息以 “data” 命名的 blob 存入 workspace 当中. 同样重要的是要记住我们添加到模型中的每个保存的网络中包含着什么. predict_net 包含着模型的结构, 包括 forward pass 的相关 ops. 而 init_net 包含着 predict_net 中的 ops 所需的参数的权重信息. 举例来说, 如果在 predict_net 中存在一个名为 fc1 的 op, 那么 init_net 就会包含这个全连接层的权重矩阵和偏置项 (fc1_w) 和 (fc1_b).
当我们 append 网络以后, 我们添加了一个 accuracy 层到模型中, 它是利用网络输出的 softmax 以及数据的真实标签来计算 accuracy 的. 最终, 我们可以利用 workspace.FetchBlob() 来手动的获取指定 blob 的值.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 用 init net 来填充(populate) model helper 对象, 用于给模型提供权重参数
init_net_proto = caffe2_pb2.NetDef()
with open(INIT_NET, "rb") as f:
    init_net_proto.ParseFromString(f.read())
test_model.param_init_net = test_model.param_init_net.AppendNet(
    core.Net(init_net_proto)
)

# 用 predict net 来添加模型的 net, 用于定义模型的结构
predict_net_proto = caffe2_pb2.NetDef()
with open(PREDICT_NET, "rb") as f:
    predict_net_proto.ParseFromString(f.read())
test_model.net = test_model.net.AppendNet(core.Net(predict_net_proto))

# 添加 accuracy, 三个参数分别为: model, blobs_in, blobs_out
accuracy = brew.accuracy(test_model, ['softmax', 'label'], 'accuracy')

接下来, 我们来运行 testing model 以查看结果.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# 先利用 init_net 初始化 test_model, 然后再创建该模型
workspace.RunNetOnce(test_model.param_init_net)
workspace.RunNet(test.net, overwrite=True)

# Stat keeper
avg_accuracy = 0.0

# 运行网络的总次数, 因为我们的测试数据集总共有 10000 张图片, 并且模型 batch 为 1, 故设置为 10000
test_iters = 10000

# Main testing loop
for i in range(test_iters):
    workspace.RunNet(test_model.net)
    acc = workspace.FetchBlob('accuracy')
    avg_accuracy += acc
    if (i % 500 == 0) and (i > 0):
        print("Iter:{}, Current Accuracy: {}".format(i, avg_accuracy/float(i)))

# 输出相关信息
print("*************************")
print("Final Test Accuracy: ", avg_accuracy/float(test_iters))

输出结果如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
Iter: 500, Current Accuracy: 0.654
Iter: 1000, Current Accuracy: 0.663
Iter: 1500, Current Accuracy: 0.640666666667
Iter: 2000, Current Accuracy: 0.643
Iter: 2500, Current Accuracy: 0.6356
Iter: 3000, Current Accuracy: 0.639333333333
Iter: 3500, Current Accuracy: 0.639142857143
Iter: 4000, Current Accuracy: 0.63625
Iter: 4500, Current Accuracy: 0.636444444444
Iter: 5000, Current Accuracy: 0.6352
Iter: 5500, Current Accuracy: 0.638181818182
Iter: 6000, Current Accuracy: 0.638666666667
Iter: 6500, Current Accuracy: 0.638461538462
Iter: 7000, Current Accuracy: 0.636428571429
Iter: 7500, Current Accuracy: 0.634533333333
Iter: 8000, Current Accuracy: 0.63475
Iter: 8500, Current Accuracy: 0.633058823529
Iter: 9000, Current Accuracy: 0.632111111111
Iter: 9500, Current Accuracy: 0.632421052632
*********************************************
Final Test Accuracy:  0.6321

利用 checkpoint 对模型进行再训练

从上面的结果我们可以看出, 我们模型在精度上虽然比随机猜测要高一些, 但是如果能够进行更多次迭代的训练, 模型的精度应该还会进一步提高, 为此, 我们需要进行以下步骤:

  • 创建一个新的 model helper
  • 指定训练数据集和 training imdb
  • 利用 Add_Original_CIFAR10_Model() 重新定义模型的结构.
  • 从之前保存的 init_net.pb 文件中获取相应的参数权重
  • 继续训练

下面我们新创建了一个 model helper 对象用于训练, 请注意要将 init_params 参数设置为 False. 这一点非常重要, 因为我们不希望 brew 自动初始化参数, 而是从文件中加载权重(也就是我们自己来设置这些参数的值). 一旦我们创建了 model helper, 我们将会添加 input layer 并将其指向 training lmdb.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# 定义训练的迭代次数
training_iters = 3000

# 清空当前的工作空间(主要是清空之前的 testing stage 产生的信息)
workspace.ResetWorkspace()

# 创建一个新的 model
arg_scope = {"order": "NCHW"}
train_model = model_helper.ModelHelper(
    name="cifar10_train", arg_scope=arg_scope, init_params=False
)

# 向模型中添加 data layer
data, _ =  AddInputLayer(train_model, 100, TRAIN_LMDB, 'lmdb')
softmax = Add_Original_CIFAR10_Model(train_model, data, 10, 32, 32, 3)

# 用之前保存的文件来填充 model 对象的 param_init_net 成员
init_net_proto = caffe2_pb2.NetDef()
with open(INIT_NET, "rb") as f:
    init_net_proto.ParseFromString(f.read())
tmp_init_net = core.Net(init_net_proto)
train_model.param_init_net = train_model.param_init_net.AppendNet(tmp_init_net)

下面的代码指定了损失函数和优化器, 和之前的代码没什么不同.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 向模型中添加 training operators
xent = train_model.LabelCrossEntropy([softmax, 'label'], 'xent')
# 计算 loss
loss = train_model.AveragedLoss(xent, "loss")
# 根据模型的精度
accuracy = brew.accuracy(train_model, [softmax, "label"], "accuracy")
# 添加计算梯度的 ops
train_model.AddGradientOperators([loss])
# 指定优化算法
optimizer.build_ssd(
    train_model,
    base_learning_rate=0.01,
    policy="fixed",
    momentum=0.9,
    weight_decay=0.004
)

注意:
我们可以利用 GetOptimizationParamInfo() 函数来获得会被优化函数优化的参数. 如果你希望以不同的方式训练你的模型, 而模型看起来并没有被正确训练, 那么检查一下这个函数的返回值. 如果返回值为空, 那么就说明定义的网络没有训练任何东西! 这就是我们在 Add_Original_CIFAR10_Model() 函数中利用 brew 来创建网络层的原因之一, 因为 brew 会自动帮我们创建网络层所需的参数. 而如果我们已经在 Model Helper 中 append .net, 那么就会返回空, 意味着没有任何参数被优化. 在工作区中, 当你 append 一个 net 以后, 你需要利用 create_param() 手动的创建 params.

1
2
for param in train_model.GetOptimizationParamInfo():
    print("Param to be optimized: ", param)

正常情况下, 应该输出如下:

1
2
3
4
5
6
7
8
9
10
Param to be optimized:  conv3_b
Param to be optimized:  fc2_w
Param to be optimized:  fc1_b
Param to be optimized:  conv1_b
Param to be optimized:  conv2_b
Param to be optimized:  conv3_w
Param to be optimized:  fc2_b
Param to be optimized:  fc1_w
Param to be optimized:  conv2_w
Param to be optimized:  conv1_w

定义好模型以后, 输入下面的代码开始训练:

1
2
3
4
5
6
7
8
9
10
11
# 初始化并创建 train model 的网络
workspace.RunNetOnce(train_model.param_init_net)
workspace.CreateNet(train_model.net, overwrite=True)

# Run the training loop
for i in range(training_iters):
    workspace.RunNet(train_model.net)
    acc = workspace.FetchBlob('accuracy')
    loss = workspace.FetchBlob('loss')
    if i % 100 == 0:
        print("Iter: {}, Loss: {}, Accuracy: {}".format(i, loss, acc))

正常情况下, 输出如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
Iter: 0, Loss: 1.02017736435, Accuracy: 0.629999995232
Iter: 100, Loss: 1.03627979755, Accuracy: 0.639999985695
Iter: 200, Loss: 1.01044285297, Accuracy: 0.649999976158
Iter: 300, Loss: 1.11098098755, Accuracy: 0.569999992847
Iter: 400, Loss: 0.857563138008, Accuracy: 0.660000026226
Iter: 500, Loss: 0.947329521179, Accuracy: 0.689999997616
Iter: 600, Loss: 1.03314650059, Accuracy: 0.660000026226
Iter: 700, Loss: 0.916410207748, Accuracy: 0.72000002861
Iter: 800, Loss: 0.780662596226, Accuracy: 0.730000019073
Iter: 900, Loss: 0.84335398674, Accuracy: 0.709999978542
Iter: 1000, Loss: 0.813905596733, Accuracy: 0.740000009537
Iter: 1100, Loss: 0.696036875248, Accuracy: 0.730000019073
Iter: 1200, Loss: 0.901836633682, Accuracy: 0.680000007153
Iter: 1300, Loss: 0.892839670181, Accuracy: 0.649999976158
Iter: 1400, Loss: 0.832642018795, Accuracy: 0.72000002861
Iter: 1500, Loss: 0.771974146366, Accuracy: 0.72000002861
Iter: 1600, Loss: 0.926190078259, Accuracy: 0.699999988079
Iter: 1700, Loss: 0.884124755859, Accuracy: 0.680000007153
Iter: 1800, Loss: 0.891266524792, Accuracy: 0.660000026226
Iter: 1900, Loss: 0.633636891842, Accuracy: 0.75
Iter: 2000, Loss: 0.643068671227, Accuracy: 0.77999997139
Iter: 2100, Loss: 0.827954411507, Accuracy: 0.740000009537
Iter: 2200, Loss: 0.845260441303, Accuracy: 0.689999997616
Iter: 2300, Loss: 0.806422412395, Accuracy: 0.689999997616
Iter: 2400, Loss: 0.742512345314, Accuracy: 0.740000009537
Iter: 2500, Loss: 0.785759627819, Accuracy: 0.759999990463
Iter: 2600, Loss: 0.530824780464, Accuracy: 0.790000021458
Iter: 2700, Loss: 0.726455926895, Accuracy: 0.709999978542
Iter: 2800, Loss: 0.945818066597, Accuracy: 0.670000016689
Iter: 2900, Loss: 0.67063999176, Accuracy: 0.810000002384

利用 Retrained 模型进行推演

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
arg_scope = {"order": "NCHW"}
# 构建模型
test_model = model_helper.ModelHelper(
    name="test_model", arg_scope=arg_scope, init_params=False
)
# 将输入层指向 test lmdb
data, _ = AddInputLayer(test_model, 1, TEST_LMDB, 'lmdb')
# 创建模型结构
softmax = Add_Original_CIFAR10_Model(test_model, data, 10, 32, 32, 3)
accuracy = brew.accuracy(test_model, ['softmax', 'label'], 'accuracy')

# Prime the net (准备网络)
workspace.RunNetOnce(test_model.param_init_net)
workspace.CreateNet(test_model.net, overwrite=True)

# CIFAR10 的混合矩阵
cmat = np.zeros((10, 10))

# Stat keepers
avg_accuracy = 0.0
test_iters = 10000

# Main testing loop
for i in range(test_iters):
    workspace.RunNet(test_model.net)
    acc = workspace.FetchBlob('accuracy')
    avg_accuracy += acc
    if (i % 500 == 0) and (i > 0):
        print("Iter: {}, Current Accuracy: {}".format(i, avg_accuracy/float(i)))

    # 获取 top-1 预测结果
    results = workspace.FetchBlob('softmax')[0]
    label = workspace.FetchBlob('label')[0]
    max_index, max_value = max(enumerate(results), key=operator.itemgetter(1))

    # 更新混合矩阵
    cmat[label, max_index] += 1

# Report final testing results
print("***************************")
print("Final Test Accuracy: ", avg_accuracy/float(test_iters))

正常情况下, 输出如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
Iter: 500, Current Accuracy: 0.712
Iter: 1000, Current Accuracy: 0.721
Iter: 1500, Current Accuracy: 0.718
Iter: 2000, Current Accuracy: 0.713
Iter: 2500, Current Accuracy: 0.71
Iter: 3000, Current Accuracy: 0.716
Iter: 3500, Current Accuracy: 0.72
Iter: 4000, Current Accuracy: 0.71675
Iter: 4500, Current Accuracy: 0.719555555556
Iter: 5000, Current Accuracy: 0.7194
Iter: 5500, Current Accuracy: 0.720363636364
Iter: 6000, Current Accuracy: 0.721166666667
Iter: 6500, Current Accuracy: 0.720461538462
Iter: 7000, Current Accuracy: 0.719571428571
Iter: 7500, Current Accuracy: 0.7188
Iter: 8000, Current Accuracy: 0.7185
Iter: 8500, Current Accuracy: 0.718235294118
Iter: 9000, Current Accuracy: 0.718
Iter: 9500, Current Accuracy: 0.717157894737
*********************************************
Final Test Accuracy:  0.7169

查看结果

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# Plot confusion matrix
fig = plt.figure(figsize=(10,10))
plt.tight_layout()
ax = fig.add_subplot(111)
res = ax.imshow(cmat, cmap=plt.cm.rainbow,interpolation='nearest')
width, height = cmat.shape
for x in xrange(width):
    for y in xrange(height):
        ax.annotate(str(cmat[x,y]), xy=(y, x),horizontalalignment='center',verticalalignment='center')

classes = ['Airplane','Automobile','Bird','Cat','Deer','Dog','Frog','Horse','Ship','Truck']
plt.xticks(range(width), classes, rotation=0)
plt.yticks(range(height), classes, rotation=0)
ax.set_xlabel('Predicted Class')
ax.set_ylabel('True Class')
plt.title('CIFAR-10 Confusion Matrix')
plt.show()

输出如下: