理解神经网络更好：第5篇——代码助手对决

在之前的系列博客文章[1, 2, 3, 4]中，我进行了一些实验，绘制了全连接泄漏ReLU网络在训练重现输入图像时生成的多面体边界。当我尝试将实验扩展到更大的网络时，我注意到代码出现了显著的减速，这是由于在CPU上计算激活模式的哈希值导致的——因此每个训练步骤都会很快，但随后所有内容都会因为可视化而陷入停滞，对于每个像素，代码都会前向评估神经网络（总共1024*1024次），每当预测被计算时，它会将激活模式传输到CPU，然后执行哈希计算。这非常慢，而且非常不并行。

我曾考虑编写一些自定义的CUDA代码来加速——没有理由存储激活模式或传输它，解决这个问题的“正确”方法是在飞行中计算哈希，理想情况下是一个具有可交换更新函数的哈希，这样不同ReLU神经元更新哈希的顺序就不重要了。

然而，这是一个业余项目，我目前没有时间做任何过于聪明的事情。因此，我决定——在做任何复杂的事情之前——看看我是否可以让我经常使用的两个现有编码助手之一为我解决这个问题。

因此，我创建了两个不同的目录，将相同的基础仓库检出到两者中，在两个目录中创建分支，然后使用以下提示查询Gemini CLI和Claude Code执行任务：

此目录中的Python脚本训练一个全连接泄漏ReLU网络以重现输入图像。它还绘制图片说明ReLU在输入空间中创建的折痕生成的多面体边界。不幸的是，生成多面体可视化的代码很慢，因为它涉及1024*1024次神经网络前向评估，然后需要将激活模式哈希成一个哈希值，以唯一标识像素所在的多面体。

我想通过将哈希计算嵌入到GPU上的前向传递中来加速此计算，而不是在最后计算激活模式的哈希。这可能可以通过PyTorch钩子实现，但我不确切知道。

我知道的是，如果我运行
1
python3 ./draw-poly-while-training.py --input ./centered_ring.png --shape [100]*20 --epochs 30 --seed 12345678 --points 5050 --save-interval 10
输出看起来像这样：
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
(...)
Input size (MB): 0.01
Forward/backward pass size (MB): 16.39
Params size (MB): 0.77
Estimated Total Size (MB): 17.17
==========================================================================================
2025-07-08 15:15:25,811 - polytope_nn - INFO - Epoch 1/2000000 - Train Loss: 3.315190, Val Loss: 0.329414
2025-07-08 15:15:25,857 - polytope_nn - INFO - Epoch 2/2000000 - Train Loss: 1.045730, Val Loss: 0.065818
2025-07-08 15:15:25,901 - polytope_nn - INFO - Epoch 3/2000000 - Train Loss: 1.414065, Val Loss: 0.488735
2025-07-08 15:15:25,948 - polytope_nn - INFO - Epoch 4/2000000 - Train Loss: 0.201550, Val Loss: 0.102159
2025-07-08 15:15:26,100 - polytope_nn - INFO - Epoch 5/2000000 - Train Loss: 0.198983, Val Loss: 0.050712
2025-07-08 15:15:26,145 - polytope_nn - INFO - Epoch 6/2000000 - Train Loss: 0.255710, Val Loss: 0.060731
2025-07-08 15:15:26,189 - polytope_nn - INFO - Epoch 7/2000000 - Train Loss: 0.122960, Val Loss: 0.091274
2025-07-08 15:15:26,232 - polytope_nn - INFO - Epoch 8/2000000 - Train Loss: 0.180629, Val Loss: 0.053913
2025-07-08 15:15:26,276 - polytope_nn - INFO - Epoch 9/2000000 - Train Loss: 0.826762, Val Loss: 0.156673
2025-07-08 15:15:26,320 - polytope_nn - INFO - Epoch 10/2000000 - Train Loss: 0.211313, Val Loss: 0.117810
2025-07-08 15:16:27,853 - polytope_nn - INFO - Visualization @ epoch 10: 61.53s
2025-07-08 15:16:27,899 - polytope_nn - INFO - Epoch 11/2000000 - Train Loss: 0.174978, Val Loss: 0.053103
2025-07-08 15:16:27,943 - polytope_nn - INFO - Epoch 12/2000000 - Train Loss: 0.332561, Val Loss: 0.095801
2025-07-08 15:16:27,987 - polytope_nn - INFO - Epoch 13/2000000 - Train Loss: 0.192859, Val Loss: 0.064341
2025-07-08 15:16:28,031 - polytope_nn - INFO - Epoch 14/2000000 - Train Loss: 0.115424, Val Loss: 0.051763
2025-07-08 15:16:28,076 - polytope_nn - INFO - Epoch 15/2000000 - Train Loss: 0.362009, Val Loss: 0.128609
2025-07-08 15:16:28,122 - polytope_nn - INFO - Epoch 16/2000000 - Train Loss: 0.117143, Val Loss: 0.058641
2025-07-08 15:16:28,165 - polytope_nn - INFO - Epoch 17/2000000 - Train Loss: 0.335812, Val Loss: 0.082517
2025-07-08 15:16:28,211 - polytope_nn - INFO - Epoch 18/2000000 - Train Loss: 0.079342, Val Loss: 0.060753
2025-07-08 15:16:28,257 - polytope_nn - INFO - Epoch 19/2000000 - Train Loss: 0.104123, Val Loss: 0.047914
2025-07-08 15:16:28,304 - polytope_nn - INFO - Epoch 20/2000000 - Train Loss: 0.097466, Val Loss: 0.050452
2025-07-08 15:17:31,553 - polytope_nn - INFO - Visualization @ epoch 20: 63.25s
从中我们可以看到，对于这种大小的网络，单个可视化步骤需要超过一分钟，分析显示大部分时间花在CPU上的哈希计算，而不是GPU。我希望你找到一种方法，在GPU上的前向传递期间计算哈希，理想情况下不将激活向量存储在内存中，而是有一个可以可交换更新的哈希函数，这样每个ReLU单元可以在计算前向传递时更新最终哈希。

我希望你：

创建一个合理的计划来改进和加速代码。

实施该计划。

使用指定的命令行重新运行脚本，并观察是否确实发生了加速——例如检查（a）可视化是否加速和（b）10个训练步骤和可视化一起是否加速。

加速可视化步骤但减慢训练步骤以至于10个训练步骤和1个可视化步骤变得更慢是非常容易的。

还请验证更改前和更改后版本的图像输出是否相同，以确保更改不会破坏任何内容。

然后我让两个模型运行了一段时间。两个模型都提供了更改，但Gemini未能实际验证结果是否相同。Claude一次就解决了问题；Gemini需要以下额外的提示：

我已经运行了你的示例代码，并检查了输出。更改前和更改后版本的输出图像不相同，甚至训练损失也发生了变化。值得一提的是，在你的版本中看不到任何多面体。你能重新检查你的工作，这次确保你检查输出是否相同吗？

有了这个额外的推动/提示，模型提供的解决方案完美无缺，甚至比Claude版本快一点点。

让我们看看两个模型生成的代码：Gemini分支和Claude分支。阅读更改，一些事情变得清晰：

Gemini在RNG上搬起石头砸自己的脚，生成了一堆随机哈希系数，这搞乱了RNG的状态，因此训练运行在更改前后不再具有可比性。
Gemini使用torch.matmul进行哈希计算，而Claude将哈希计算为torch.sum(A * B)。
Claude将代码分解为更多的小函数，而Gemini没有。Claude的代码稍微更可读，Gemini的更改更最小。

有趣的东西。两个解决方案都不是我完全想要的，但它们目前足够好，并且比我开始时（也是 vibe-coded）的东西提供了相当显著的加速。这是我第一次编码助手以非平凡的方式帮助我优化代码，这……当然是一些东西。

无论如何，通过这些优化，我现在可以在稍大的具有数百万参数的神经网络上运行我的数据可视化电影生成，所以还有更多的学习在前面。我现在需要弄清楚如何以编程方式上传YouTube视频，但与此同时，这里有一个视频，训练一个100神经元、10层深的网络，执行我之前文章中的“圆圈绘制”任务。Vibe coding随机改变了我的线条颜色，但嘿，没关系。

像往常一样，这个视频中的问题比答案多。最让我困惑的是后期训练中训练的相对“不稳定性”。这在“闪烁”中可见，似乎随机地SGD步骤遇到 vastly higher loss，部分屏幕变黑，损失飙升，然后训练需要恢复。有趣的是，在这些情况下，多面体的几何形状变化不大，但许多多面体上的线性函数同时变化，以一种对整体性能非常不利的方式。一旦程序化上传工作，我将上传更多视频，因为我有一个有趣的观察：

当训练发散时（对于更大更深的网络），发散首先通过搞乱线性函数开始，只有在它们被彻底搞乱之后，多面体的几何形状也开始变得混乱。

直到那时！

神经网络优化实战：通过代码助手加速GPU哈希计算

本文详细介绍了如何利用代码助手优化神经网络训练中的哈希计算，通过将激活模式哈希计算嵌入GPU前向传播，显著提升可视化步骤性能，同时比较了Gemini和Claude两种代码助手的实现方案。

理解神经网络更好：第5篇——代码助手对决