- 收集每个GPU上的输出
在分布式训练时，每个GPU都会有一部分数据，当我们需要使用全部的数据进行计算时，我们需要收集所有GPU的tensor。
比如两个GPU，第一个GPU有16组数据，第二个GPU有16组数据， 在进行对比学习计算时，我们需要收集所有的输出来增加负样本的数量。
我们可以使用tensors_from_all = self.all_gather(my_tensor)
比如：
def training_step(self, batch, batch_idx):
outputs = self(batch)
...
all_outputs = self.all_gather(outputs, sync_grads=True)
loss = contrastive_loss_fn(all_outputs, ...)
return loss
最后编辑于 ：
©著作权归作者所有,转载或内容合作请联系作者
【社区内容提示】社区部分内容疑似由AI辅助生成，浏览时请结合常识与多方信息审慎甄别。
平台声明：文章内容（如有图片或视频亦包括在内）由作者上传并发布，文章内容仅代表作者本人观点，简书系信息发布平台，仅提供信息存储服务。
