@Team
2019-06-08T14:07:11.000000Z
字数 11566
阅读 4866
叶虎
TensorFlow 2.0全面keras化:如果你想使用高级的layers,只能选择keras。TensorFlow 1.x存在tf.layers以及tf.contrib.slim等高级API来创建模型,但是2.0仅仅支持tf.keras.layers,不管怎么样,省的大家重复造轮子,也意味着模型构建的部分大家都是统一的,增加代码的复用性(回忆一下原来的TensorFlow模型构建真是千奇百怪)。值得注意的tf.nn模块依然存在,里面是各种常用的nn算子,不过大部分人不会去直接用这些算子构建模型,因为keras.layers基本上包含了常用的网络层。当然,如果想构建新的layer,可以直接继承tf.keras.layers.Layer:
class Linear(tf.keras.layers.Layer):
def __init__(self, units=32, **kwargs):
super(Linear, self).__init__(**kwargs)
self.units = units
def build(self, input_shape):
self.w = self.add_weight(shape=(input_shape[-1], self.units),
initializer='random_normal',
trainable=True)
self.b = self.add_weight(shape=(self.units,),
initializer='random_normal',
trainable=True)
def call(self, inputs):
return tf.matmul(inputs, self.w) + self.b
layer = Linear(32)
print(layer.weights) # [] the weights have not created
x = tf.ones((8, 16))
y = layer(x) # shape [8, 32]
print(layer.weights)
这里我们继承了Layer来实现自定义layer。第一个要注意的点是我们定义了build方法,其主要用于根据input_shape创建layer的Variables。注意,我们没有在类构造函数中创建Variables,而是单独定义了一个方法。之所以这样做类的构造函数中并没有传入输入Tensor的信息,这里需要的是input的输入特征维度,所以无法创建Variables。这个build方法会在layer第一次真正执行(执行layer(input))时才会执行,并且只会执行一次(Layer内部有self.build这个bool属性)。这是一种懒惰执行机制,如果熟悉Pytorch的话,PyTorch在创建layer时是需要输入Tensor的信息,这意味着它是立即创建了Variables。
第二点是Layer本身有很多属性和方,这里列出一些重要的:
Layer类是keras中最基本的类,对其有个全面的认识比较重要,具体可以看源码。大部分情况下,我们只会复用keras已有的layers,而我们创建模型最常用的是keras.Model类,这个Model类是继承了Layer类,但是提供了更多的API,如model.compile(), model.fit(), model.evaluate(), model.predict()等,熟悉keras的都知道这是用于模型训练,评估和预测的方法。另外重要的一点,我们可以继承Model类,创建包含多layers的模块或者模型:
class ConvBlock(tf.keras.Model):
"""Convolutional Block consisting of (conv->bn->relu).
Arguments:
num_filters: number of filters passed to a convolutional layer.
kernel_size: the size of convolution kernel
weight_decay: weight decay
dropout_rate: dropout rate.
"""
def __init__(self, num_filters, kernel_size,
weight_decay=1e-4, dropout_rate=0.):
super(ConvBlock, self).__init__()
self.conv = tf.keras.layers.Conv2D(num_filters,
kernel_size,
padding="same",
use_bias=False,
kernel_initializer="he_normal",
kernel_regularizer=tf.keras.regularizers.l2(weight_decay))
self.bn = tf.keras.layers.BatchNormalization()
self.dropout = tf.keras.layers.Dropout(dropout_rate)
def call(self, x, training=True):
output = self.conv(x)
output = self.bn(x, training=training)
output = tf.nn.relu(output)
output = self.dropout(output, training=training)
return output
model = ConvBlock(32, 3, 1e-4, 0.5)
x = tf.ones((4, 224, 224, 3))
y = model(x)
print(model.layers)
这里我们构建了一个包含Conv2D->BatchNorm->ReLU的block,打印model.layers可以获得其内部包含的所有layers。更进一步地,我们可以在复用这些block就像使用tf.keras.layers一样构建更复杂的模块:
class SimpleCNN(tf.keras.Model):
def __init__(self, num_classes):
super(SimpleCNN, self).__init__()
self.block1 = ConvBlock(16, 3)
self.block2 = ConvBlock(32, 3)
self.block3 = ConvBlock(64, 3)
self.global_pool = tf.keras.layers.GlobalAveragePooling2D()
self.classifier = tf.keras.layers.Dense(num_classes)
def call(self, x, training=True):
output = self.block1(x, training=training)
output = self.block2(output, training=training)
output = self.block3(output, training=training)
output = self.global_pool(output)
logits = self.classifier(output)
return logits
model = SimpleCNN(10)
print(model.layers)
x = tf.ones((4, 32, 32, 3))
y = model(x) # [4, 10]
这种使用手法和PyTorch的Module是类似的,并且Model类的大部分属性会递归地收集内部layers的属性,比如model.weights是模型内所有layers中定义的weights。
构建模型的另外方式还可以采用Keras原有方式,如采用tf.keras.Sequential:
model = tf.keras.Sequential([
# Adds a densely-connected layer with 64 units to the model:
layers.Dense(64, activation='relu', input_shape=(32,)),
# Add another:
layers.Dense(64, activation='relu'),
# Add a softmax layer with 10 output units:
layers.Dense(10, activation='softmax')])
或者采用keras的functional API:
inputs = keras.Input(shape=(784,), name='img')
x = layers.Dense(64, activation='relu')(inputs)
x = layers.Dense(64, activation='relu')(x)
outputs = layers.Dense(10, activation='softmax')(x)
model = keras.Model(inputs=inputs, outputs=outputs, name='mnist_model')
虽然都可以,但是我个人还是喜欢第一种那种模块化的模型构建方法。另外,你可以对call方法应用tf.function,这样模型执行就使用Graph模式了。
在开始模型训练之前,一个重要的项是数据加载,TensorFlow 2.0的数据加载还是采用tf.data,不过在eager模式下,tf.data.Dataset这个类将成为一个Python迭代器,我们可以直接取值:
dataset = tf.data.Dataset.range(10)
for i, elem in enumerate(dataset):
print(elem) # prints 0, 1, ..., 9
这里我们只是展示了一个简单的例子,但是足以说明tf.data在TensorFlow 2.0下的变化,tf.data其它使用技巧和TensorFlow 1.x是一致的。
另外tf.keras提供两个重要的模块losses和metrics用于模型训练。对于losses,其本身就是对各种loss函数的封装,如下面的case:
bce = tf.keras.losses.BinaryCrossentropy()
loss = bce([0., 0., 1., 1.], [1., 1., 1., 0.])
print('Loss: ', loss.numpy()) # Loss: 11.522857
而metrics模块主要包含了常用的模型评估指标,这个模块与TensorFlow 1.x的metrics模块设计理念是一致的,就是metric本身是有状态的,一般是通过创建Variable来记录。基本用法如下:
m = tf.keras.metrics.Accuracy()
m.update_state([1, 2, 3, 4], [0, 2, 3, 4])
print('result: ', m.result().numpy()) # result: 0.75
m.update_state([0, 2, 3], [1, 2, 3])
print('result: ', m.result().numpy()) # result: 0.714
m.reset_states() # 重置
m.update_state([0, 2, 3], [1, 2, 3])
print('result: ', m.result().numpy()) # result: 0.667
当你需要自定义metric时,你可以继承tf.keras.metrics.Metric类,然后实现一些接口即可,下面这个例子展示如何计算多分类问题中TP数量:
class CatgoricalTruePositives(tf.keras.metrics.Metric):
def __init__(self, name='categorical_true_positives', **kwargs):
super(CatgoricalTruePositives, self).__init__(name=name, **kwargs)
self.true_positives = self.add_weight(name='tp', initializer='zeros')
def update_state(self, y_true, y_pred, sample_weight=None):
y_pred = tf.argmax(y_pred)
values = tf.equal(tf.cast(y_true, 'int32'), tf.cast(y_pred, 'int32'))
values = tf.cast(values, 'float32')
if sample_weight is not None:
sample_weight = tf.cast(sample_weight, 'float32')
values = tf.multiply(values, sample_weight)
self.true_positives.assign_add(tf.reduce_sum(values))
def result(self):
return self.true_positives
def reset_states(self):
# The state of the metric will be reset at the start of each epoch.
self.true_positives.assign(0.)
上面的三个接口必须都要实现,其中update_state是通过添加新数据而更新状态,而reset_states是重置初始值,result方法是获得当前状态,即metric结果。注意这个metric其实是创建了一个Variable来保存TP值。你可以类比实现更复杂的metric。
对于模型训练,我们可以通过下面一个完整实例来全面学习:
import numpy as np
import tensorflow as tf
fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
# Adding a dimension to the array -> new shape == (28, 28, 1)
train_images = train_images[..., None]
test_images = test_images[..., None]
# Getting the images in [0, 1] range.
train_images = train_images / np.float32(255)
test_images = test_images / np.float32(255)
train_labels = train_labels.astype('int64')
test_labels = test_labels.astype('int64')
# dataset
train_ds = tf.data.Dataset.from_tensor_slices(
(train_images, train_labels)).shuffle(10000).batch(32)
test_ds = tf.data.Dataset.from_tensor_slices(
(test_images, test_labels)).batch(32)
# Model
class MyModel(tf.keras.Sequential):
def __init__(self):
super(MyModel, self).__init__([
tf.keras.layers.Conv2D(32, 3, activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Conv2D(64, 3, activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10, activation=None)
])
model = MyModel()
# optimizer
initial_learning_rate = 1e-4
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
initial_learning_rate,
decay_steps=100000,
decay_rate=0.96,
staircase=True)
optimizer = tf.keras.optimizers.RMSprop(learning_rate=lr_schedule)
# checkpoint
checkpoint = tf.train.Checkpoint(step=tf.Variable(0), optimizer=optimizer, model=model)
manager = tf.train.CheckpointManager(checkpoint, './tf_ckpts', max_to_keep=3)
# loss function
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
# metric
train_loss_metric = tf.keras.metrics.Mean(name='train_loss')
train_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
test_loss_metric = tf.keras.metrics.Mean(name='test_loss')
test_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')
# define a train step
@tf.function
def train_step(inputs, targets):
with tf.GradientTape() as tape:
predictions = model(inputs, training=True)
loss = loss_object(targets, predictions)
loss += sum(model.losses) # add other losses
# compute gradients and update variables
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_loss_metric(loss)
train_acc_metric(targets, predictions)
# define a test step
@tf.function
def test_step(inputs, targets):
predictions = model(inputs, training=False)
loss = loss_object(targets, predictions)
test_loss_metric(loss)
test_acc_metric(targets, predictions)
# train loop
epochs = 10
for epoch in range(epochs):
print('Start of epoch %d' % (epoch,))
# Iterate over the batches of the dataset
for step, (inputs, targets) in enumerate(train_ds):
train_step(inputs, targets)
checkpoint.step.assign_add(1)
# log every 20 step
if step % 20 == 0:
manager.save() # save checkpoint
print('Epoch: {}, Step: {}, Train Loss: {}, Train Accuracy: {}'.format(
epoch, step, train_loss_metric.result().numpy(),
train_acc_metric.result().numpy())
)
train_loss_metric.reset_states()
train_acc_metric.reset_states()
# do test
for inputs, targets in test_ds:
test_step(inputs, targets)
print('Test Loss: {}, Test Accuracy: {}'.format(
test_loss_metric.result().numpy(),
test_acc_metric.result().numpy()))
麻雀虽小,但五脏俱全,这个实例包括数据加载,模型创建,以及模型训练和测试。特别注意的是,这里我们将train和test的一个step通过tf.function转为Graph模式,可以加快训练速度,这是一种值得推荐的方式。另外一点,上面的训练方式采用的是custom training loops,自由度较高,另外一种训练方式是采用keras比较常规的compile和fit训练方式。
TensorFlow 2.0的另外一个特点是提供tf.distribute.Strategy
更好地支持分布式训练,其接口更加简单易用。我们最常用的分布式策略是单机多卡同步训练,tf.distribute.MirroredStrategy
完美支持这种策略。这种策略将在每个GPU设备上创建一个模型副本(replica),模型中的参数在所有replica之间映射,称之为MirroredVariables,当他们执行相同更新时将在所有设备间同步。底层的通信采用all-reduce算法,all-reduce方法可以将多个设备上的Tensors聚合在每个设备上,这种通信方式比较高效,而all-reduce算法有多中实现方式,这里默认采用NVIDIA NCCL的all-reduce方法。创建这种策略只需要简单地定义:
mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"],
cross_device_ops=tf.distribute.NcclAllReduce())
# 这里将在GPU 0和1上同步训练
当我们创建好分布式策略后,在后续的操作中只需要加入strategy.scope即可。下面我们创建一个简单的模型以及优化器:
with mirrored_strategy.scope():
model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
optimizer = tf.keras.optimizers.SGD(learning_rate=0.001)
对于dataset,我们需要调用tf.distribute.Strategy.experimental_distribute_dataset
来分发数据:
with mirrored_strategy.scope():
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(1000).batch(
global_batch_size)
# 注意这里是全局batch size
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)
然后我们定义train step,并采用strategy.experimental_run_v2
来执行:
@tf.function
def train_step(dist_inputs):
def step_fn(inputs):
features, labels = inputs
with tf.GradientTape() as tape:
logits = model(features)
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(
logits=logits, labels=labels)
loss = tf.reduce_sum(cross_entropy) * (1.0 / global_batch_size)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
return cross_entropy
per_example_losses = mirrored_strategy.experimental_run_v2(step_fn, args=(dist_inputs,))
mean_loss = mirrored_strategy.reduce(tf.distribute.ReduceOp.MEAN,
per_example_losses, axis=0)
return mean_loss
这里要注意的是我们要将loss除以全部batch size,只是因为分布式训练时在更新梯度前会将所有replica上梯度通过all-reduce算法相加聚合到每个设备上。另外,strategy.experimental_run_v2
返回是每个replica的结果,要得到最终结果,需要reduce聚合一下。
最后是执行训练,采用循环方式即可:
with mirrored_strategy.scope():
for inputs in dist_dataset:
print(train_step(inputs))
要注意的是MirroredStrategy只支持单机多卡同步训练,如果想使用多机版本,需要采用MultiWorkerMirorredStrateg。其它的分布式训练策略还有CentralStorageStrategy,TPUStrategy,ParameterServerStrategy。想深入了解的话,可以查看distribute_strategy guide以及distribute_strategy tuorial。
这里我们简明扼要地介绍了TensorFlow 2.0的核心新特性,相信掌握这些新特性就可以快速入手TensorFlow 2.0。不过目前Google只发布了TensorFlow 2.0.0-beta0版本,未来也许会有更多想象不到的黑科技。加油!TensorFlow Coders。