resnet-18 网络 训练cifar100数据-tf2.0-详细注释

resnet.py

import tensorflow as tf
from   tensorflow import keras
from   tensorflow.keras import layers, Sequential



class BasicBlock(layers.Layer):
    # 继承自layers.layer 这个父类。实现2个基本函数：__init__ 初始化函数  call 函数

    def __init__(self, filter_num, stride=1):
        # __init__ 中定义网络层，做初始化。call中使用__init__中定义好的网络，进行前向传播（计算，得到输出）
        # filter_num ：通道数量 贯穿始终。  stride：=1时，如果做padding 则图片大小不变。如果stride !=1 则会对输入做padding，使其可以对stride整除。比如stride=2，输入为32*32，则输出为16*16

        super(BasicBlock, self).__init__()

        self.conv1 = layers.Conv2D(filter_num, (3, 3), strides=stride, padding="same")
        # 选用小的卷积核，准确度不会下降，所以一般选择 1*1 3*3 这样的卷积核 两个卷积层中，第一个卷积层strides=stride 有可能执行下采样 第二个stride=1,即第二个卷积层肯定不会执行下采样，这样调用BasicBlock时，即使stride不为1，下采样只会执行一次。
        self.bn1 = layers.BatchNormalization()
        self.relu = layers.Activation('relu')

        self.conv2 = layers.Conv2D(filter_num, (3, 3), strides=1, padding='same')
        #                                              第二个stride=1
        self.bn2 = layers.BatchNormalization()

        if stride != 1:
            self.downsample = Sequential()
            self.downsample.add(layers.Conv2D(filter_num, (1, 1), strides=stride))
            #     .add :向Sequential中添加层       通过1*1的卷积核 和strides改变其形状。使其与通过卷积层后的数据形状相同。
        else:
            self.downsample = lambda x:x
            # lambda：             参数：函数体。举例：b=lambda x: x+1 相当于 def g(x):return x+1

    def call(self, inputs, training=None):
        # call函数有2个基本参数 inputs 和 training。 这是两个函数的标准的写法

        # [b, h, w, c]
        out = self.conv1(inputs)
        out = self.bn1(out, training=training)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out, training=training)

        identify = self.downsample(inputs)

        output = layers.add([out, identify])
        # 两个输出直接做相加，得到output  layers下面有一个add，把这2个层添加进来相加。
        output = tf.nn.relu(output)

        return output

# Res Block 模块。继承keras.Model或者keras.Layer都可以
class ResNet(keras.Model):

    # 第一个参数layer_dims：[2, 2, 2, 2] 4个Res Block，每个包含2个Basic Block
    # 第二个参数num_classes：我们的全连接输出，取决于输出有多少类。
    def __init__(self, layer_dims, num_classes=100):

        super(ResNet, self).__init__()

        # resnet-18的第一层:预处理层。实现起来比较灵活可以加 MAXPool2D，可以没有。
        self.stem = Sequential([layers.Conv2D(64, (3,3), strides=(1, 1)),
                               layers.BatchNormalization(),
                               layers.Activation('relu'),
                               layers.MaxPool2D(pool_size=(2, 2),strides=(1, 1), padding='same')
                                ])

        # 创建4个Res Block；注意第1项不一定以2倍形式扩张，都是比较随意的，这里都是经验值。
        self.layers1 = self.build_resblock(64, layer_dims[0])
        self.layers2 = self.build_resblock(128, layer_dims[1], stride=2)
        self.layers3 = self.build_resblock(256, layer_dims[2], stride=2)
        self.layers4 = self.build_resblock(512, layer_dims[3], stride=2)

        # 残差网络输出output: [b, 512, h, w];长宽无法确定，上面的需要运算一下，如果这里没有办法确定的话。
        # 用这个层可以自适应的确定输出。表示不管你的长和宽是多少，我会在某一个channel上面，所有的长和宽像素值加起来
        # 求一个均值，比如：有512个3*3的feature map,[512, 3, 3],每个feature map为3*3,9个像素值，我做一个这样的
        # average，得到一个平均的像素值是多少。下面这里处理之后得到一个512的vector，准确来说为[512, 1, 1],这个512的
        # vector就可以送到先形成进行分类。
        # [b, 512, h, w]  GlobalAveragePooling2D 作用是把一个feature map 上的像素值做一个均值 输出是[b, 512]
        self.avgpool = layers.GlobalAveragePooling2D()
        # 全连接层 输出类别是num_classes
        self.fc = layers.Dense(num_classes)

    def call(self, inputs, training=None):
        # __init__中准备工作完毕；下面完成前向运算过程。
        x = self.stem(inputs)

        x = self.layers1(x)
        x = self.layers2(x)
        x = self.layers3(x)
        x = self.layers4(x)

        # 做一个global average pooling，得到之后只会得到一个channel，不需要做reshape操作了。
        # shape为 [batchsize, channel]
        # [b, 512]
        x = self.avgpool(x)
        # [b, 100]
        x = self.fc(x)

        return x

    # 实现 Res Block； 创建一个Res Block  filter_num：通道数量    blocks： block数量，即堆叠的BasicBlock数量。
    def build_resblock(self, filter_num, blocks, stride=1):

        res_blocks = Sequential()
        # 添加第一个BasicBlock
        # may down sample 也许进行下采样。
        # 对于当前Res Block中的Basic Block，我们要求每个Res Block只有一次下采样的能力。
        res_blocks.add(BasicBlock(filter_num, stride))
        # 添加后续的BasicBlock 设置stride=1 不让它执行下采样。
        for _ in range(1, blocks):
            res_blocks.add(BasicBlock(filter_num, stride=1))  # 这里stride设置为1，只会在第一个Basic Block做一个下采样。

        return res_blocks


def resnet18():

    return ResNet([2, 2, 2, 2])
# 如果我们要使用 ResNet-34 的话，那34是怎样的配置呢？只需要改一下这里就可以了。对于56，152去查一下配置
def resnet34():

    return ResNet([3, 4, 6, 3])  #4个Res Block，第1个包含3个Basic Block,第2为4，第3为6，第4为3




# 补充：下面需要使用的。
#
# 介绍一下global average pooling ，这个概念出自于 network in network；global average pooling 与 average pooling 的差别就在 “global” 这一个字眼上。global 与 local 在字面上都是用来形容 pooling 窗口区域的。 local 是取 feature map 的一个子区域求平均值，然后滑动这个子区域； global 显然就是对整个 feature map 求平均值了。
# 主要是用来解决全连接的问题，其主要是是将最后一层的特征图进行整张图的一个均值池化，形成一个特征点，将这些特征点组成最后的特征向量进行softmax中进行计算
# 举个例子：假如，最后的一层的数据是10个6×6的特征图，global average pooling是将每一张特征图计算所有像素点的均值，输出一个数据值。这样10 个特征图就会输出10个数据点，将这些数据点组成一个1×10的向量的话，就成为一个特征向量，就可以送入到softmax的分类中计算了;

# 如果在training时出现out of memory,就是显卡的显存比较小网络参数比较大，可以适当调小batchsize
# 可以尝试修改resnet。标准的是有4个resblock 每个包含2个basicblock，可以把它改成[1, 1, 1, 1] 这样每个resnet 包含一个basicblock，参数大概减少了一半
# 可以尝试google的colab，提供了一个k80的显卡，11G的显存。

resnet-18-train.py

import os
import tensorflow as tf
from tensorflow.keras import layers, optimizers, datasets, Sequential
from resnet import resnet18

os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
tf.random.set_seed(2345)

# 数据预处理，仅仅是类型的转换。    [-1~1]
def preprocess(x, y):

    x = 2 * tf.cast(x, dtype=tf.float32) / 255. - 1.
    y = tf.cast(y, dtype=tf.int32)

    return x, y

# 加载cifar100数据 如果本地没有，会自动下载。
(x, y), (x_test, y_test) = datasets.cifar100.load_data()
# print(x.shape, y.shape, x_test.shape, y_test.shape)
# 形状：(50000, 32, 32, 3) (50000, 1)这里要注意，y的形状,后面要去掉为1的维度  (10000, 32, 32, 3) (10000, 1)

y = tf.squeeze(y, axis=1)           # 或者tf.squeeze(y, axis=1)把1维度的squeeze掉。 去掉为1 的维度 [50000, 1] => [5000]
y_test = tf.squeeze(y_test, axis=1) # 去掉为1 的维度 [50000, 1] => [5000]

# 加载数据集
train_db = tf.data.Dataset.from_tensor_slices((x, y))
# 对数据的处理       打乱数据        批预处理         设置batch大小
train_db = train_db.shuffle(10000).map(preprocess).batch(32)

test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_db = test_db.map(preprocess).batch(32)

# 这里只是为了查看数据形状，在整个程序中并无意义
sample = next(iter(train_db))
print("sample:", sample[0].shape, sample[1].shape, tf.reduce_min(sample[0]), tf.reduce_max(sample[0]))



def main():

    # 输入[b, 32, 32, 3]
    model = resnet18()
    # build()
    model.build(input_shape=(None, 32, 32, 3))
    model.summary() # 查看网络结构和参数量
    # 优化器：Adam
    optimizer = optimizers.Adam(lr=1e-3)

    # 训练流程：
    for epoch in range(50):

        for step, (x, y) in enumerate(train_db):

            # 从数据进入网络，到计算损失函数，都要放到with tf.GradientTape() as tape: 中，方便计算梯度
            with tf.GradientTape() as tape:

                # [b, 32, 32, 3] => [b, 100]
                logits = model(x)
                # 对y做one_hot处理 [b] => [b, 100]
                y_one_hot = tf.one_hot(y, depth=100)
                # 计算loss 结果维度[b]   交叉熵损失函数      one_hot之后的y  输出结果  这里要设置为True
                loss = tf.losses.categorical_crossentropy(y_one_hot, logits, from_logits=True)
                loss = tf.reduce_mean(loss)

            # 计算梯度  variables：所有变量
            grads = tape.gradient(loss, model.trainable_variables)
            # 梯度更新     grads 要和 变量一一对应，所以用zip
            optimizer.apply_gradients(zip(grads, model.trainable_variables))

            if step % 100 == 0:
                print(epoch, step, "loss:", float(loss))


        # 以下是测试部分--
        # 测试流程：
        # 1、数据放入网络得到数据。
        # 2、 softmax得到概率。
        # 3、 argmax 得到最大值的位置。
        # 4、equal 与最大值比较， 得到True False。 cast转化为int格式（0， 1）
        # 5、reduce_sum 计算预测正确的数量
        # 6、计算总的预测正确的数量（每个batch预测正确的相加） 和 总的数据数量（每个batch的总数量）
        # 7、计算准确率 acc
        total_num = 0
        total_correct = 0
        for x, y in test_db:

            # 数据放入网络，得到输出 [b, 100]
            logits = model(x, training=False)

            # softmax 处理 每个数据处理成概率
            prob = tf.nn.softmax(logits, axis=1)
            # argmax 返回指定维度的最大值的位置 int64类型
            pred = tf.argmax(prob, axis=1)
            # 转化数据格式 int64 -> int32
            pred = tf.cast(pred, dtype=tf.int32)

            # tf.cast：把true False转化成 int格式：0或者1   tf.equal: 比较预测值和真实值，结果是True或者False
            correct = tf.cast(tf.equal(pred, y), dtype=tf.int32)
            # 计算所有正确的总数（在上一步，预测正确的已经转化为1）
            correct = tf.reduce_sum(correct)

            # 总数量，是把每一批的数量加起来
            total_num += x.shape[0]
            # 预测正确的总数量，把每次预测正确的总数量加起来
            total_correct += int(correct)

        # 计算准确率
        acc = total_correct / total_num
        print(epoch, 'acc:', acc)



if __name__ == '__main__':
    main()