编程速记(41)：PyTorch篇&Tensorflow篇-CUDA out of memory & 指定model/data运行GPU

一、简述

场景简述：文件需要加载多个不同的模型同时运行从而导致cuda OOM，这些模型有的是用tf代码写的，有的是用PyTorch写的。

二、尝试过的失败方案

对于PyTorch而言，使用多GPU训练方案也许是一种可行的方案。

net = torch.nn.DataParallel(model, device_ids=[xxx])

但是由于也涉及到tf的model，因此这样并不能解决问题，运行起来依旧会报CUDA oom。

三、解决方案

1.调小batch_size

可以尝试将batch_size直接降到1，如果还是不行就只能考虑第二种方案。

2.手动指定不同model放到不同的GPU上

对于tf：

with tf.device('/gpu:2'):
    com, rec = ComCNN(), RecCNN()
    com.summary()
    rec.summary()

对于PyTorch：

    model = model.to(torch.device('cuda: X'))

或者是：

    model = model.cuda(X) # X代表GPU编号

值得注意的是：不同的tensor之间如果需要进行运算的话，需要保证在同一个gpu上，将tensor指定到具体的gpu上：

    data = data.cuda(X) # X代表GPU编号

四、参考

1.torch.device的作用
2.Torch之CUDA使用指定显卡