Pytorch переместит nn.module в cuda, включая подмодули и модули-оболочки ⇐ Python

[Anonymous]

Я новичок в использовании Pytorch. Я хотел бы переместить свой модуль в cuda.
Пример кода:

Код: Выделить всё

class SubModel(nn.Module):
def __init__(self):
super(SubModel, self).__init__()
self.conv1 = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=2)

def forward(self, x):
print(f"x type:{type(x)}")
print(f"weight type:{type(self.conv1.weight)}")
return self.conv1(x)

class WrapperModel(nn.Module):
def __init__(self, count):
super(WrapperModel, self).__init__()
self.blocks = []
for i in range(count):
self.blocks.append(SubModel())

def forward(self, x):
for block in self.blocks:
x = block(x)
return x

class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.conv = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=2)
self.wrapper = WrapperModel(2)

def forward(self, x):
x = self.conv(x)
x = self.wrapper(x)

return x

Однако я столкнулся с ошибкой времени выполнения при выполнении self.conv1(x) в прямой функции в SubModel:
RuntimeError: Тип ввода (torch.cuda.FloatTensor ) и тип веса (torch.FloatTensor) должны быть одинаковыми
Вывод команд печати:

Код: Выделить всё

x type:
weight type:

Затем я попытался переместить обертку в cuda в конце:

Код: Выделить всё

model.wrapper.to(device)

это не помогло.
Наконец, я перемещаю conv1D из подмодели в cuda:

Код: Выделить всё

class SubModel(nn.Module):
def __init__(self):
super(SubModel, self).__init__()
self.conv1 = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=2).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))

Все работает отлично! Однако результат печати не изменился:

Код: Выделить всё

x type:
weight type:

Поэтому я понятия не имею, что здесь не так.
Я также чувствую, что слишком неудобно вручную перемещать все экземпляры в cuda.Мой вопрос: существует ли удобный способ переместить модель в cuda, включая все ее подмодели? Спасибо!

Подробнее здесь: https://stackoverflow.com/questions/786 ... er-modules