我必须用不同的设备在不同种类的pytorch模型上堆叠一些我自己的层。
例如,A是一个cuda模型,B是一个cpu模型(但在得到设备类型之前我并不知道它)。那么新的模型分别是C和D,其中
class NewModule(torch.nn.Module):
def __init__(self, base):
super(NewModule, self).__init__()
self.base = base
self.extra = my_layer() # e.g. torch.nn.Linear()
def forward(self,x):
y = self.base(x)
z = self.extra(y)
return z
...
C = NewModule(A) # cuda
D = NewModule(B) # cpu
但是我必须将base
和extra
移到相同的设备上,也就是说C的base
和extra
是cuda模型,D的是cpu模型。所以我尝试了__inin__
:
def __init__(self, base):
super(NewModule, self).__init__()
self.base = base
self.extra = my_layer().to(base.device)
不幸的是,torch.nn.Module
中没有属性device
(引发AttributeError
)。
我应该怎么做才能获得base
的设备类型?或者使用任何其他方法使base
和extra
自动位于同一个设备上,即使base
的结构是不特定的?
4条答案
按热度按时间vh0rcniy1#
这个问题已经被问过很多次了(1,2)。引用一个PyTorch开发人员的回答:
That’s not possible. Modules can hold parameters of different types on different devices, and so it’s not always possible to unambiguously determine the device.
推荐的工作流程(如PyTorch博客中所述)是单独创建
device
对象并在任何地方使用它。请注意,没有什么可以阻止您向模型添加
.device
属性。正如Kani(在注解中)提到的,如果模型中的所有参数都在同一个设备上,则可以使用
next(model.parameters()).device
。h9vpoimq2#
我的解决方案,在99%的情况下有效。
此后,dummy_param将始终与模块Net具有相同的设备,因此您可以随时获取它。例如:
00jrzges3#
@Duane的答案在模型中创建了一个参数(尽管是一个小Tensor)。
我觉得这个答案稍微有点蟒和雅:
net.cuda()
、net.float()
等都将同样工作,因为它们都调用_apply
而不是to
(如在the source中所见)。来自@Kani的评论的替代解决方案(接受的答案)也非常优雅:
对于参数,您通过
model.device
访问设备。当模型中没有参数时,此解决方案不起作用。e0uiprwp4#
您可以使用
next(model.parameters()).device
。