网站首页>/ 文章专栏>/ 深度学习框架对比：PyTorch vs TensorFlow
而在众多深度学习框架中，PyTorch 和 TensorFlow 无疑是两颗最为耀眼的明星，它们各自凭借独特的魅力，吸引着无数开发者与研究者投身其中。那么，这两者究竟谁更胜一筹？是 PyTorch 的简洁灵活，还是 TensorFlow 的强大稳健？接下来，就让我们一同深入剖析，探寻答案。
PyTorch 的诞生极富传奇色彩，它最初是由 Facebook 人工智能研究院（FAIR）团队精心打造的。追溯其根源，可回溯到 2002 年诞生的 Torch，Torch 采用了 Lua 语言作为接口，虽然在当时的科研领域得到了一定的应用，但 Lua 语言的小众性限制了其更广泛的传播。随着 Python 在科学计算领域的日益普及，2017 年，Torch 的幕后团队毅然对 Tensor 之上的所有模块进行了全面重构，推出了 PyTorch。这一举措犹如在深度学习的天空中点亮了一颗璀璨的新星，PyTorch 迅速在 GitHub 上崭露头角，一举占据热度榜榜首 。
TensorFlow 的起源同样不凡，它最初是由 Google Brain 团队开发，基于 Google 自主研发的深度学习基础架构 DistBelief 构建而来。2015 年 11 月 9 日，Google 依据 Apache 2.0 协议将其开源，这一开源之举在深度学习领域掀起了巨大的波澜，无数开发者得以接触和使用这个强大的工具，TensorFlow 也借此迅速在全球范围内传播开来，成为了机器学习领域的重要力量。
相比之下，PyTorch 在内存管理上具有一定的相对优势。PyTorch 采用了一种更加灵活的内存分配策略，它的动态计算图是随着代码的执行逐步构建的，只有在需要时才会分配内存，并且在计算完成后，会及时释放不再使用的内存空间。这种按需分配和及时释放的机制，使得 PyTorch 在处理小规模和中等规模数据集时，内存占用相对较低。此外，PyTorch 还提供了一些内存优化的工具和方法，如梯度累积等，进一步提高了内存使用效率。在处理大规模数据集时，通过合理地使用这些工具和方法，PyTorch 也能够有效地减少内存占用，避免内存不足的问题。
TensorFlow 则通过对静态计算图的优化，来提高执行速度。如前面提到的，TensorFlow 可以在运行前对计算图进行剪枝、节点合并等优化操作，减少计算量和中间数据传输开销，从而提高执行速度。此外，TensorFlow 还支持使用 XLA（Accelerated Linear Algebra）编译器，对计算图进行进一步的优化，将计算图编译成高效的机器码，提高执行效率。在 TensorFlow 2.0 版本中，引入了 Eager Execution 模式，使得 TensorFlow 在保留静态图性能优势的同时，也具备了动态图的灵活性，进一步提升了执行速度和开发效率。
import torchimport torch.nn as nnimport torch.optim as optimfrom torch.utils.data import DataLoaderfrom torchvision import datasets, transforms# 数据预处理transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])# 加载训练集和测试集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)# 定义神经网络模型class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 10, kernel_size=5) self.conv2 = nn.Conv2d(10, 20, kernel_size=5) self.fc1 = nn.Linear(320, 50) self.fc2 = nn.Linear(50, 10) def forward(self, x): x = nn.functional.relu(nn.functional.max_pool2d(self.conv1(x), 2)) x = nn.functional.relu(nn.functional.max_pool2d(self.conv2(x), 2)) x = x.view(-1, 320) x = nn.functional.relu(self.fc1(x)) x = self.fc2(x) return nn.functional.log_softmax(x, dim=1)# 初始化模型、损失函数和优化器model = Net()criterion = nn.NLLLoss()optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)# 训练模型def train(epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() if batch_idx % 100 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item()))# 测试模型def test(): model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: output = model(data) test_loss += criterion(output, target).item() pred = output.data.max(1, keepdim=True)[1] correct += pred.eq(target.data.view_as(pred)).sum() test_loss /= len(test_loader.dataset) print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset)))# 开始训练和测试for epoch in range(1, 10 + 1): train(epoch) test()import tensorflow as tffrom tensorflow.keras import datasets, layers, models# 加载并预处理数据(train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()train_images, test_images = train_images / 255.0, test_images / 255.0# 增加通道维度，因为卷积层需要四维输入(batch_size, height, width, channels)train_images = train_images[..., tf.newaxis]test_images = test_images[..., tf.newaxis]# 构建模型model = models.Sequential([ layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)), layers.MaxPooling2D((2, 2)), layers.Conv2D(64, (3, 3), activation='relu'), layers.MaxPooling2D((2, 2)), layers.Flatten(), layers.Dense(64, activation='relu'), layers.Dense(10)])# 编译模型model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy'])# 训练模型model.fit(train_images, train_labels, epochs=5)# 评估模型test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=2)print('\nTest accuracy:', test_acc)# 预测probability_model = tf.keras.Sequential([model, tf.keras.layers.Softmax()])predictions = probability_model.predict(test_images)