TensorFlow深度学习框架详解

TensorFlow是Google开发的开源深度学习框架，提供完整的工具生态系统，支持从研究原型到生产部署的全流程。

tf.keras核心API

tf.keras是TensorFlow的高级API，提供简洁的模型构建方式。Sequential模型适合线性层堆叠，Functional API支持复杂网络拓扑。

import tensorflow as tf
from tensorflow import keras

# Sequential模型
model = keras.Sequential([
    keras.layers.Input(shape=(784,)),
    keras.layers.Dense(256, activation='relu'),
    keras.layers.Dropout(0.3),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dropout(0.3),
    keras.layers.Dense(10, activation='softmax')
])

model.summary()

# Functional API构建复杂模型
inputs = keras.Input(shape=(32, 32, 3))
x = keras.layers.Conv2D(32, (3, 3), activation='relu')(inputs)
x = keras.layers.MaxPooling2D((2, 2))(x)
x = keras.layers.Conv2D(64, (3, 3), activation='relu')(x)
x = keras.layers.GlobalAveragePooling2D()(x)
outputs = keras.layers.Dense(10, activation='softmax')(x)

cnn_model = keras.Model(inputs=inputs, outputs=outputs)

模型编译与训练

编译时配置优化器、损失函数和指标。fit方法支持批处理、验证集和回调函数。

model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.001),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 准备数据
import numpy as np
X_train = np.random.randn(2000, 784).astype('float32')
y_train = np.random.randint(0, 10, (2000,))

# 训练模型
history = model.fit(
    X_train, y_train,
    epochs=15,
    batch_size=64,
    validation_split=0.2,
    callbacks=[
        keras.callbacks.EarlyStopping(patience=3, restore_best_weights=True),
        keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=2)
    ]
)

自定义训练循环

使用GradientTape实现完全自定义的训练逻辑，适合需要精细控制的场景。

optimizer = keras.optimizers.Adam(learning_rate=0.001)
loss_fn = keras.losses.SparseCategoricalCrossentropy()

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        loss = loss_fn(y, predictions)
        loss += sum(model.losses)
    
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss, predictions

for epoch in range(10):
    epoch_loss = 0.0
    for i in range(0, len(X_train), 64):
        x_batch = tf.constant(X_train[i:i+64])
        y_batch = tf.constant(y_train[i:i+64])
        loss, _ = train_step(x_batch, y_batch)
        epoch_loss += loss.numpy()
    print(f"Epoch {epoch+1}, Loss: {epoch_loss/(len(X_train)//64):.4f}")

分布式训练

TensorFlow通过分布式策略实现多GPU和多机训练，MirroredStrategy支持单机多卡数据并行。

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    model = keras.Sequential([
        keras.layers.Dense(256, activation='relu', input_shape=(784,)),
        keras.layers.BatchNormalization(),
        keras.layers.Dense(128, activation='relu'),
        keras.layers.Dense(10, activation='softmax')
    ])
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

print(f"可用设备: {strategy.num_replicas_in_sync} GPUs")

模型保存与部署

支持多种保存格式：SavedModel用于生产部署，Keras格式便于继续训练，TFLite用于移动端。

# 保存完整模型
model.save('saved_model/my_model')

# 保存权重
model.save_weights('model_weights.h5')

# 加载模型
loaded_model = keras.models.load_model('saved_model/my_model')

# TFLite转换
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

print(f"TFLite模型大小: {len(tflite_model)/1024:.1f} KB")

TensorFlow凭借完善的生态系统和工业级部署能力，成为深度学习框架的主流选择之一。