Python深度学习，TensorFlow与PyTorch核心用法与实战对比

一、TensorFlow与PyTorch前言

深度学习已成为计算机视觉、自然语言处理、推荐系统等领域的核心技术，而TensorFlow和PyTorch是当前最主流的两大Python深度学习框架。本文从零基础入门角度，系统讲解两大框架的核心概念、基础用法及实战场景，帮助初学者快速掌握框架使用逻辑，避开入门常见坑。

二、框架选型：TensorFlow vs PyTorch怎么选？

1. 适用场景对比

特性	TensorFlow	PyTorch
开发风格	静态图为主，2.x版本支持动态图（Eager Execution）	动态图优先，调试直观，代码更贴近Python原生逻辑
应用场景	工业部署、大规模分布式训练、移动端/嵌入式部署	学术研究、快速原型开发、小批量实验迭代
生态完善度	官方工具链丰富（TensorBoard、TFLite、TF Serving）	第三方库适配多（HuggingFace、TorchVision）
学习曲线	稍陡，概念抽象（如Session、Graph）	平缓，Python开发者易上手

2. 选型建议

• 若你是入门新手/科研人员：优先选PyTorch，动态图调试方便，代码可读性高，易理解深度学习底层逻辑；
• 若你是工业开发/部署需求：优先选TensorFlow，完善的部署工具链能覆盖从训练到上线的全流程；
• 若你无明确方向：先学PyTorch建立基础认知，再补充TensorFlow核心概念，两大框架底层逻辑相通，触类旁通。

三、环境搭建：统一安装流程（避免版本兼容问题）

1. 前置条件

• Python 3.8~3.10（兼容两大框架的稳定版本）；
• 显卡配置（可选）：NVIDIA显卡+CUDA 11.7+cuDNN 8.4（无显卡则用CPU版本）。

2. 安装命令

（1）PyTorch安装（推荐官网匹配命令）

# CPU版本
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2
# GPU版本（适配CUDA 11.7）
pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu117

（2）TensorFlow安装

# CPU版本
pip install tensorflow==2.13.0
# GPU版本
pip install tensorflow==2.13.0 tensorflow-gpu==2.13.0

（3）验证安装

# PyTorch验证
import torch
print("PyTorch版本：", torch.__version__)
print("是否支持CUDA：", torch.cuda.is_available())  # 输出True则GPU可用

# TensorFlow验证
import tensorflow as tf
print("TensorFlow版本：", tf.__version__)
print("是否支持CUDA：", tf.config.list_physical_devices('GPU'))  # 输出GPU列表则可用

四、核心概念与基础实战（以MNIST手写数字识别为例）

1. 核心概念统一理解

无论TensorFlow还是PyTorch，深度学习的核心流程一致：

数据加载 → 模型定义 → 损失函数与优化器 → 训练循环 → 模型评估

以下通过MNIST实战，对比两大框架的实现逻辑，帮你快速掌握核心用法。

2. PyTorch实战：MNIST手写数字识别

步骤1：数据加载与预处理

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 数据预处理：归一化+转换为张量
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图片转为Tensor，像素值0-255转为0-1
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差
])

# 加载训练集和测试集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 数据加载器（批量处理+打乱数据）
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

步骤2：定义神经网络模型

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # 卷积层1：输入1通道（灰度图），输出16通道，卷积核3x3
        self.conv1 = nn.Conv2d(1, 16, 3, 1, 1)
        # 池化层：2x2最大池化
        self.pool = nn.MaxPool2d(2, 2)
        # 卷积层2：输入16通道，输出32通道
        self.conv2 = nn.Conv2d(16, 32, 3, 1, 1)
        # 全连接层1：32*7*7（池化后尺寸）→ 128
        self.fc1 = nn.Linear(32 * 7 * 7, 128)
        # 全连接层2：128 → 10（10个数字类别）
        self.fc2 = nn.Linear(128, 10)
        # 激活函数
        self.relu = nn.ReLU()

    def forward(self, x):
        # 前向传播：卷积→激活→池化
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        # 展平张量：(batch_size, 32, 7, 7) → (batch_size, 32*7*7)
        x = x.view(-1, 32 * 7 * 7)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 实例化模型
model = SimpleCNN()
# 若有GPU则移到GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

步骤3：定义损失函数与优化器

# 损失函数：交叉熵损失（适用于分类任务）
criterion = nn.CrossEntropyLoss()
# 优化器：Adam优化器，学习率0.001
optimizer = optim.Adam(model.parameters(), lr=0.001)

步骤4：训练循环

epochs = 3  # 训练轮数
for epoch in range(epochs):
    running_loss = 0.0
    model.train()  # 切换到训练模式
    for i, (images, labels) in enumerate(train_loader):
        # 将数据移到指定设备（CPU/GPU）
        images, labels = images.to(device), labels.to(device)
        
        # 梯度清零（避免累加）
        optimizer.zero_grad()
        # 前向传播
        outputs = model(images)
        # 计算损失
        loss = criterion(outputs, labels)
        # 反向传播
        loss.backward()
        # 更新参数
        optimizer.step()
        
        running_loss += loss.item()
        # 每100批次打印一次损失
        if i % 100 == 99:
            print(f'Epoch [{epoch+1}/{epochs}], Batch [{i+1}], Loss: {running_loss/100:.4f}')
            running_loss = 0.0
print('训练完成！')

步骤5：模型评估

correct = 0
total = 0
model.eval()  # 切换到评估模式（禁用Dropout/BatchNorm等）
with torch.no_grad():  # 禁用梯度计算，节省内存
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        # 获取预测结果（概率最大的类别）
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'测试集准确率: {100 * correct / total:.2f}%')
# 保存模型
torch.save(model.state_dict(), 'mnist_cnn.pth')

3. TensorFlow实战：MNIST手写数字识别

步骤1：数据加载与预处理

import tensorflow as tf
from tensorflow.keras import datasets, layers, models

# 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = datasets.mnist.load_data()
# 归一化：像素值0-255转为0-1，并增加通道维度（适配卷积层）
x_train = x_train.reshape((60000, 28, 28, 1)).astype('float32') / 255
x_test = x_test.reshape((10000, 28, 28, 1)).astype('float32') / 255
# 标签独热编码（可选，也可直接用稀疏交叉熵）
y_train = tf.keras.utils.to_categorical(y_train, 10)
y_test = tf.keras.utils.to_categorical(y_test, 10)

步骤2：定义神经网络模型

model = models.Sequential([
    # 卷积层1：32个3x3卷积核，激活函数ReLU，输入尺寸28x28x1
    layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    # 池化层：2x2最大池化
    layers.MaxPooling2D((2, 2)),
    # 卷积层2：64个3x3卷积核
    layers.Conv2D(64, (3, 3), activation='relu'),
    layers.MaxPooling2D((2, 2)),
    # 卷积层3：64个3x3卷积核
    layers.Conv2D(64, (3, 3), activation='relu'),
    # 展平层
    layers.Flatten(),
    # 全连接层1：64个神经元
    layers.Dense(64, activation='relu'),
    # 全连接层2：10个神经元，softmax激活（输出概率）
    layers.Dense(10, activation='softmax')
])

# 打印模型结构
model.summary()

步骤3：编译模型（损失函数+优化器+评估指标）

model.compile(optimizer='adam',
              loss='categorical_crossentropy',  # 独热编码标签用该损失
              # loss='sparse_categorical_crossentropy',  # 非独热编码标签用该损失
              metrics=['accuracy'])

步骤4：训练模型

# 训练模型：批量大小64，训练3轮，验证集比例0.1
history = model.fit(x_train, y_train,
                    batch_size=64,
                    epochs=3,
                    validation_split=0.1)

步骤5：模型评估与保存

# 评估测试集
test_loss, test_acc = model.evaluate(x_test, y_test)
print(f'测试集准确率: {test_acc:.4f}')

# 保存模型
model.save('mnist_cnn_tf.h5')

五、入门常见问题与解决方案

1. 数据维度不匹配

• 问题表现：报错“Input 0 of layer conv2d is incompatible with the layer”；
• 解决方案：检查输入数据的维度，确保与模型输入层shape一致（如PyTorch是[batch, channel, height, width]，TensorFlow是[batch, height, width, channel]）。

2. GPU内存不足

• 问题表现：PyTorch报错“CUDA out of memory”，TensorFlow报错“Resource exhausted”；
• 解决方案：减小batch_size（如从64改为32）、使用梯度累积、关闭不必要的程序释放GPU内存。

3. 过拟合

• 问题表现：训练集准确率高，测试集准确率低；
• 解决方案：添加Dropout层（PyTorch：nn.Dropout(0.2)；TensorFlow：layers.Dropout(0.2)）、数据增强、减少模型参数。