Tensorflow 2.0 搭建网络练习

使用Tensorflow 2.0 实现Resnet-50 ，根据TF官方示例改写，分别用model.fit() 和tf.function 实现模型训练

Tensorflow 2.0 对keras进行了集成，兼顾了易用性和灵活性，之前只使用过Keras搭建模型，这里结合tensorflow 和keras实现Resnet-50。供学习交流使用，希望大佬们指出问题！！

---

以下是代码实现（编辑器jupyter notebook）：

导入 fashion mnist 数据集

from tensorflow.keras.datasets import fashion_mnist
import tensorflow as tf
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
(x_train, y_train), (x_test, y_test) = fashion_mnist.load_data()
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

y_train = tf.keras.utils.to_categorical(y_train)
y_test = tf.keras.utils.to_categorical(y_test)
print(x_train.shape)
print(y_train.shape)

查看可用GPU数量

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

定义网络输入时的卷积层，用了3个3x3的卷积核代替原来的 7x7的卷积

class Preprocess_Block(tf.keras.Model):
    def __init__(self):
        super(Preprocess_Block, self).__init__()
        self.conv_1 = tf.keras.layers.Conv2D(64, kernel_size=3, padding='same', strides=1, data_format='channels_last')
        self.conv_2 = tf.keras.layers.Conv2D(64, kernel_size=3, padding='same', strides=1, data_format='channels_last')
        self.conv_3 = tf.keras.layers.Conv2D(64, kernel_size=3, padding='same', strides=1, data_format='channels_last')
        self.max_pooling = tf.keras.layers.MaxPooling2D(pool_size=2, strides=1, padding='same')
    def call(self, x, training=None):
        x = self.conv_1(x)
        x = self.conv_2(x)
        x = self.conv_3(x)
        y = self.max_pooling(x)
        return y

定义网络的基本块

class Bottleneck_Block(tf.keras.Model):
    def __init__(self, filters, stride=1, is_upsampling=False):
        super(Bottleneck_Block, self).__init__()
        filter_num_1,filter_num_2,filter_num_3 = filters
        self.is_upsampling = is_upsampling
        if self.is_upsampling:
            self.up_sampling = tf.keras.layers.Conv2D(filter_num_3, kernel_size=1, strides=stride, data_format='channels_last')
            
        self.conv_1 = tf.keras.layers.Conv2D(filter_num_1, kernel_size=1, strides=1, data_format='channels_last')
        self.bn_1 = tf.keras.layers.BatchNormalization()
        
        self.conv_2 = tf.keras.layers.Conv2D(filter_num_2,kernel_size=3, padding='same',strides=stride, data_format='channels_last')
        self.bn_2 = tf.keras.layers.BatchNormalization()
        
        self.conv_3 = tf.keras.layers.Conv2D(filter_num_3, kernel_size=1,padding= 'same', strides=1, data_format='channels_last')
        self.bn_3 = tf.keras.layers.BatchNormalization()
        
        self.relu = tf.keras.layers.Activation('relu')
        
    def call(self, x, training=None):
        
        if self.is_upsampling:
            origin = self.up_sampling(x)
        else:
            origin = x
        x = self.conv_1(x)
        x = self.bn_1(x)
        x = self.relu(x)
        x = self.conv_2(x)
        x = self.bn_2(x)
        x = self.relu(x)
        x = self.conv_3(x)
        x = self.bn_3(x)
        x += origin
        outputs = tf.nn.relu(x)
        return outputs

定义大的 block

class Res_Block(tf.keras.Model):
    def __init__(self,filter_num, stride=1, block_number=1):
        super(Res_Block, self).__init__()
        self.block = tf.keras.models.Sequential()
        self.block.add(Bottleneck_Block(filters=filter_num, stride=stride, is_upsampling=True))
        for i in range(block_number-1):
            self.block.add(Bottleneck_Block(filters=filter_num, stride=1, is_upsampling=False))
    def call(self, x, training=None):
        output = self.block(x)
        return output

用上述定义的基本网络块，搭建ResNet-50

class ResNet(tf.keras.Model):
    def __init__(self, num_classes):
        super(ResNet, self).__init__()
        self.preprocess_block = Preprocess_Block()
        self.block_1 = Res_Block(filter_num=(64, 64, 256), stride=1, block_number=3)
        self.block_2 = Res_Block(filter_num=(128, 128, 512), stride=1, block_number=4)
        self.block_3 = Res_Block(filter_num=(256, 256, 1024), stride=1, block_number=6)
        self.block_4 =  Res_Block(filter_num=(512, 512, 2048), stride=1, block_number=6)
        self.avg_pooling_layer = tf.keras.layers.GlobalAveragePooling2D()
        self.fc_2048 = tf.keras.layers.Dense(2048, activation='relu')
        self.fc = tf.keras.layers.Dense(num_classes, activation='softmax')
        
    def call(self, inputs, training=None):
        x = self.preprocess_block(inputs)
        x = self.block_1(x)
        x = self.block_2(x)
        x = self.block_3(x)
        x = self.block_4(x)
        x = self.avg_pooling_layer(x)
        x = self.fc_2048(x)
        y = self.fc(x)
        return y

定义数据预处理程序，完成数据归一化和数据类型的转换

def preprocess(x, y):
    x = tf.cast(x, dtype=tf.float32)/255
    y = tf.cast(y, dtype=tf.uint32)
    return x, y

训练和测试数据的 生成器

train_db = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).map(preprocess).batch(10)
test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test)).map(preprocess).batch(10)

初始化模型

model = ResNet(10)

自定义训练过程

定义损失函数、优化器、度量标准

loss_function = tf.keras.losses.CategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.CategoricalAccuracy(name='train_accuracy')
test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.CategoricalAccuracy(name='test_accuracy')

自定义训练和测试过程

@tf.function
def train_per_step(images, labels):
    with tf.GradientTape() as tape:
        predictions = model(images)
        loss = loss_function(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    train_loss(loss)
    train_accuracy(labels, predictions)

@tf.function
def test_per_step(images, labels):
    predictions = model(images)
    t_loss = loss_function(labels, predictions)

    test_loss(t_loss)
    test_accuracy(labels, predictions)

import numpy as np
with tf.device('/GPU:0'):
    EPOCHS = 1
    for epoch in range(EPOCHS):
      # 在下一个epoch开始时，重置评估指标
        train_loss.reset_states()
        train_accuracy.reset_states()
        test_loss.reset_states()
        test_accuracy.reset_states()
        step = np.ceil(x_train.shape[0] / 32)
        while(step>0):
            data = next(iter(train_db))
            train_per_step(data[0], data[1])
            step -= 1
        test_step = np.ceil(x_test.shape[0] / 32)
        while(test_step > 0):
            test_data = next(iter(test_db))
            test_per_step(test_data[0], test_data[1])
            test_step -= 1
        template = 'Epoch {}, Train Loss: {},Train Accuracy: {}, Test Loss: {}, Test Accuracy: {}'
        print (template.format(epoch+1, train_loss.result(),train_accuracy.result()*100, test_loss.result(), test_accuracy.result()*100))

使用keras的model.fit()训练

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['categorical_accuracy'])
tb = tf.keras.callbacks.TensorBoard(log_dir='./log', histogram_freq=1)
model.fit_generator(train_db, callbacks=[tb], epochs=1, verbose=1)

---

后续更新 >>>> 使用Tensorflow 2.0 实现多GPU并行训练