Pytorch Lightning 分布式训练 + Taichi

综合讨论区
,
#1

请问Taichi似乎不支持新建指定GPU上的实例？
我将基于taichi实现的可微渲染器用作网络训练的Loss计算，在使用Pytorch Lightning进行分布式训练时，会报错：
[E 07/10/23 01:45:54.590 1739881] [cuda_driver.h:operator()@92] CUDA Error CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered while calling stream_synchronize (cuStreamSynchronize)

#2

这似乎是因为taichi的ti.init默认是新建实例到gpu 0，但分布式训练时，传入的tensor可能是其他gpu，请问我该怎么做呢？

#3

说不定可以使用MPI的库做一个多卡之间的通信？就是感觉通信的延迟会特别大…

#4

好像是可以指定taichi在具体显卡运行的
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
多卡的话好像这篇文章 有提过