[Pytorch使用] cuda的使用

原文

torch.cuda里，记录着当前默认device;

torch.device("cuda")，指代当前默认device；with torch.cuda.device(1)这样的region，会设置当前默认device;

torch.device('cuda:2') 是取2号GPU卡；

对input tensor进行operator操作，则output tensor也会放到和input tensor相同的device上；（无视当前默认device)

operator要求input都在相同device上，其output也会放至相同device上；跨device的operator是不存在的（tensor.to、tensor.cuda等copy操作例外）

以下是直接在GPU上开辟Tensor:

a = torch.tensor([1., 2.], device=cuda)

以下是先在Host上开辟Tensor，然后copy至GPU上：

b = torch.tensor([1., 2.]).cuda()
b2 = torch.tensor([1., 2.]).to(device=cuda)

operator如果input是在GPU上，则其操作是异步的，kernel会被发射到device的某个Stream上排队，CPU继续执行；（tensor.to等操作不同，默认是同步的；可以传参设成异步的）

可以设置环境变量CUDA_LAUNCH_BLOCKING=1，把所有operator强制变成同步的，debug时用；

使用torch.cuda.synchronize()和torch.cuda.Event来为operator发射的所有kernel计时；（只能拿到elpased time，无法拿到“时间点”）

使用torch.cuda.Stream()创建的新Stream, 和默认Stream是不同的Stream; 这些Stream之间的同步要由用户自己把控，小心冲突！

Pytorch自己对GPU显存做了一层cache；“为了避免deallocate时的同步开销”；torch.cuda.memory_summary等调用可以查看当前显存使用情况（没有哪个operator占用多少显存这些详细信息。。。）

套路：input data必须分配在或者.to复制到GPU上；model必须.to复制到GPU上；

优化：

pin_memory可以加快Host-->GPU的copy速度，可以支持异步copy（copy的参数里要设上non_blocking=True)；

DataLoader可以用参数pin_memory=True来指定把读到的数据直接放在pin_memory里面；

多卡：建议使用torch.nn.parallel.DistributedDataParallel，而不是torch.nn.DataParallel; 前者比后者快；前者是多进程，后者是多线程（受Python的GIL影响）