如何解决PyTorch:使用多个GPU推断单个非常大的图像吗?
我想在非常大的卫星图像上进行推理(即语义分割),而又不将其分割成碎片。我可以使用4个GPU(每个GPU具有15 GB的内存),并且想知道是否有可能以某种方式使用这些GPU的所有组合内存(即60 GB)来推断PyTorch中的图像?
解决方法
您正在寻找 model parallel 的工作方式。
基本上,您可以将模型的不同部分分配给不同的 GPU,然后您应该负责“簿记”。
此解决方案非常特定于模型和任务,因此没有“通用”包装器(与 data parallel 相反)。
例如:
class MyModelParallelNetwork(nn.Module):
def __init__(self,...):
# define the network
self.part_one = ... # some nn.Module
self.part_two = ... # additional nn.Module
self.part_three = ...
self.part_four = ...
# important part - "send" the different parts to different GPUs
self.part_one.to(torch.device('gpu:0'))
self.part_two.to(torch.device('gpu:1'))
self.part_three.to(torch.device('gpu:2'))
self.part_four.to(torch.device('gpu:3'))
def forward(self,x):
# forward through model parts and GPUs:
p1 = self.part_one(x.to(torch.device('gpu:0')))
p2 = self.part_two(p1.to(torch.device('gpu:1')))
p3 = self.part_three(p2.to(torch.device('gpu:2')))
y = self.part_four(p3.to(torch.device('gpu:3')))
return y # result is on cuda:3 device
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。