tensorflow2.0---笔记6 深度残差网络ResNet

ResNet

shortcut 解决网络层数加深训练艰难的问题，使得更深层网络的实现成为可能。

Basic Block

ResBlock

DenseNet

ResNet实战

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

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
tf.random.set_seed(2345)


class BasicBlock(layers.Layer):

    def __init__(self, filter_num, stride=1):
        super(BasicBlock, self).__init__()

        self.conv1 = layers.Conv2D(filter_num, (3, 3), strides=stride, padding='same')
        self.bn1 = layers.BatchNormalization()
        self.relu = layers.Activation('relu')

        self.conv2 = layers.Conv2D(filter_num, (3, 3), strides=1, padding='same')
        self.bn2 = layers.BatchNormalization()

        if stride != 1:
            self.downsample = Sequential()
            self.downsample.add(layers.Conv2D(filter_num, (1, 1), strides=stride))
        else:
            self.downsample = lambda x: x

    def call(self, inputs, training=None):

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

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

        identity = self.downsample(inputs)

        output = layers.add([out, identity])  # layers.add 实现两个层输出的内容相加
        output = tf.nn.relu(output)

        return output


class ResNet(keras.Model):

    def __init__(self, layer_dims, num_classes=100):  # [2, 2, 2, 2]
        super(ResNet, self).__init__()

        self.stem = Sequential([
            layers.Conv2D(64, (3, 3), strides=1),
            layers.BatchNormalization(),
            layers.Activation('relu'),
            layers.MaxPool2D(pool_size=(2, 2), strides=1, padding='same')
        ])

        self.layer1 = self.build_resblock(64, layer_dims[0])
        self.layer2 = self.build_resblock(128, layer_dims[1], stride=2)
        self.layer3 = self.build_resblock(256, layer_dims[2], stride=2)
        self.layer4 = self.build_resblock(512, layer_dims[3], stride=2)

        # 由于不知道输入图片大小，所以此时output：[b, 512, h, w]
        self.avgpool = layers.GlobalAveragePooling2D()  # 每个特征层取平均
        self.fc = layers.Dense(num_classes)  # 输出100类

    def build_resblock(self, filter_num, blocks, stride=1):  # 每个res_block里面有若干个Basic Block
        res_blocks = Sequential()

        # may down sample
        res_blocks.add(BasicBlock(filter_num, stride))

        for _ in range(1, blocks):
            res_blocks.add(BasicBlock(filter_num, stride=1))

        return res_blocks

    def call(self, inputs, training=None):

        x = self.stem(inputs)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = self.fc(x)

        return x


def resnet18():
    return ResNet([2, 2, 2, 2])


def resnet34():
    return ResNet([3, 4, 6, 3])


def preprocessing(x, y):

    x = tf.cast(x, dtype=tf.float32) / 255.
    y = tf.cast(y, dtype=tf.int32)
    return x, y


def main():

    # 加载数据
    (x, y), (x_test, y_test) = datasets.cifar100.load_data()
    y = tf.squeeze(y)
    y_test = tf.squeeze(y_test)

    train_db = tf.data.Dataset.from_tensor_slices((x, y))
    train_db = train_db.shuffle(1000).map(preprocessing).batch(64)

    test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test))
    test_db = test_db.map(preprocessing).batch(64)


    # 构建网络
    model = resnet18()
    model.build(input_shape=(None, 32, 32, 3))
    optimizer = optimizers.Adam(lr=1e-4)

    for epoch in range(5):

        for step, (x, y) in enumerate(train_db):
            # print(step)
            y_onehot = tf.one_hot(y, depth=100)

            with tf.GradientTape() as tape:
                logits = model(x)

                loss = tf.losses.categorical_crossentropy(y_onehot, logits, from_logits=True)
                # print(loss.shape)  # (64,)
                loss = tf.reduce_mean(loss)

            grads = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(grads, model.trainable_variables))

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

        # 每个epoch做一个测试
        total_num = 0
        correct_num = 0
        for x, y in test_db:

            logits = model(x)
            prob = tf.nn.softmax(logits, axis=1)  # 遍历每个列，即为行
            pred = tf.argmax(prob, axis=1)
            pred = tf.cast(pred, dtype=tf.int32)

            correct = tf.cast(tf.equal(pred, y), dtype=tf.int32)
            correct = tf.reduce_sum(correct)

            total_num += x.shape[0]
            correct_num += int(correct)

        acc = correct_num / total_num
        print(epoch, "test_acc:", acc)


if __name__ == '__main__':
    main()