最近在用Transformer做文本生成任务，之前一直用TensorFlow 2.x，但感觉调试起来越来越费劲，特别是自定义loss和梯度裁剪的时候，文档看得头大。朋友推荐换PyTorch，说动态图更灵活、社区也更活跃。我想问问实际用过两者的老哥，从数据加载、模型搭建到分布式训练，到底体验差别大不大？尤其是在实践层面，有没有那种“用了一个就能少掉很多坑”的情况？我目前用的是单卡V100，后面可能要上多卡。如果有比较明确的优劣对比，或者具体的踩坑经历，非常想听一听。谢谢大家！
楼主
3小时前
PyTorch和TensorFlow在Transformer训练上到底差多少？
请 登录 后发表回复
全部回复
共 2 条
2楼
2小时前
单卡V100的话，PyTorch写自定义loss和梯度裁剪确实比TF省心太多，我当初从TF切到PyTorch就因为这俩破事。数据加载方面PyTorch的DataLoader配合torchvision或者huggingface的datasets库，上手快得多。多卡训练的话，PyTorch的DistributedDataParallel基本是标配，社区踩坑经验丰富，文档也清楚，不像TF有时候得翻源码才能搞明白。建议你先用PyTorch搭个小模型跑通单卡，再往上扩，整体体验会顺滑不少。
3楼
20分钟前
从TF转PyTorch后，写自定义loss和梯度裁剪确实顺滑很多，多卡训练用DistributedDataParallel也省心，建议趁早换。