深度学习工程师私藏的代码重构三原则

　　去年二月，我蹲在三号机房调试那个混账多任务loss加权模块——就是你懂的那个，用logits做温度缩放再叠两层可学习权重的玩意儿。空调坏了，整个机房嗡嗡响得像有只蜜蜂卡在服务器风扇里，而老张坐我斜后方，一边抠指甲一边往键盘上甩薯片渣。我说这代码得动刀，他抬头看我一眼，把耳机线缠在手指上转了三圈，没说话。

　　第一原则：loss函数签名不许改。不是因为怕break backward compatibility——谁还真的care那个？是怕某天你半夜三点跑实验崩了，查日志发现居然是别人把get_loss()从(ypred, ytrue, mask)偷偷改成(ytrue, ypred, sample_weight, kwargs)，连docstring都懒得补全。我原以为只是风格问题，结果上周五下午四点十七分，CI挂了十七个pipeline，就因为小王在utils.losses.py里加了个default=1e-8的eps参数，而下游六个脚本全靠try-except硬扛。这个我敢打包票：他改完连自己写的test都没跑，光点了绿色勾。

文章配图，仅供参考

　　第二原则：DataLoader的collate_fn里不许写业务逻辑。上个月我看到有人在collate里做label smoothing——还是那种手动for循环+np.random.uniform的版本。更绝的是，他居然还在里面调了self.tokenizer.encode()，而tokenizer是挂在Dataset实例上的，等于每次collate都触发一次状态检查。我当时盯着那行代码看了十八秒，窗外正下着雨，雨滴砸在玻璃上裂成三条细痕，中间那条歪得离谱。后来我问他为啥不提出来，他说“反正batch_size=1也能跑”。我说那你把batch_size=1当训练模式？他笑了一下，端起桌上的麻辣香锅外卖盒喝了一口汤——汤上面浮着三颗花椒，全沉在底。

　　第三原则：模型forward里不许出现if config.debug == True。不为别的，就为那年七月实习生在ResNet分支里塞了个print(f"layer3_out.shape={x.shape}")，然后忘了删。上线压测那天，每步forward多打三百七十二行log，k8s日志服务直接OOM。但问题不在print，在那个if。它像一颗缓释胶囊——平时没事，等你真想关掉debug模式时，发现它早被copy-paste进五个不同repo的三个分支里，其中一个还在git blame里挂着你自己的邮箱。我翻过commit记录，最后一次修改是去年十一月凌晨一点零三分，他接了个电话，语气很急，说“妈您别动灶台”，挂了就顺手commit -am “fix bug”。那个bug其实不存在。

　　其实最烦的是optimizer.step()之后突然来个model.zero_grad()——不是位置错，是根本不需要。但大家就是爱写，跟玄学似的。我试过删，删了跑三天没出事，第四天某个teacher-student distillation的grad clip逻辑崩了，梯度norm直接飙到一百零三，我愣是没复现。所以现在我干脆留着，像供着块镇宅石。小林前两天跟我说他把zero_grad()挪到step之前，还写了benchmark，提速0.2%，我说你信吗？他摇摇头，右手食指在桌沿敲了两下，节奏跟机房空调故障声一致。

　　昨天我在办公室看见实习生用@torch.no_grad()包整个train_epoch函数。不是包val，是train。我走过去，她正啃苹果，果核上还沾着一小块红皮，她抬头问我：“是不是能省显存？”我没答。转身去倒水，听见她手机弹出消息提示音，是微信语音申请，头像是一只柴犬，耳朵被P成了蝙蝠侠造型。

　　优胜劣汰是真的。但汰掉的不全是烂代码，有时候是那个刚来三个月、写完dataloader记得测memory leak、却在PR里忘记删print(torch.cuda.memory_summary())的姑娘。她走的时候把工位擦得很干净，椅子底下漏了一颗跳跳糖，在拖地水渍里化了一半，粉红色的糖浆沿着地砖缝爬了三厘米。现在那个位置坐着新来的算法岗，昨天他推了份PR，title写着“refactor loss scaling”，commit message是“clean up”。我没点开diff。我现在只看最后两行：一行是torch.nn.functional.cross_entropy，一行是注释# keep this for warmup schedule。

　　所以啊……重构这事，不是技术问题。是你敢不敢在merge按钮按下去前，把config.yaml里那个device: cuda:1改成cuda:0试试；是你愿不愿意让QA同事把你所有loss函数的名字列张表，然后划掉其中读起来像“带温度的加权归一化KL散度补偿式混合损失”的那几个；还有——最要命的——是你能不能忍住，别在base_model.py第837行那个super().__init__()后面，加一句# fuck the embedding dim alignment bullshit。

　　我打算下周找运维要一台空GPU，重搭个干净环境，只装pytorch 2.1和numpy 1.24。不跑任何业务模型，就拿mnist+LeNet，专门测三种loss写法的backward稳定性。要是三天没崩，我就把那套“三原则”贴在茶水间微波炉旁边。要是崩了……我继续用现在这套，但把小王写的collate_fn单独抽成一个叫_legacy_collate_v2.py的文件，放进一个名叫“历史遗留荣誉墙”的目录里。名字起得越郑重，越没人敢动。

（编辑：站长网）

【声明】本站内容均来自网络，其相关言论仅代表作者个人观点，不代表本站立场。若无意侵犯到您的权利，请及时与联系站长删除相关内容!