Блог

У меня есть следующий фиктивный код в PyTorch

 class inside(nn.Module):
    def __init__(self):
        super(inside, self).__init__()        
        self.weight_h = nn.Parameter(SOMETHING GOES HERE) # send to CPU or GPU
        self.weight_v = nn.Conv2d(SOMETHING GOES HERE) # send to CPU or GPU

    def forward(self, x):
        ...
        return x

class main_class(nn.Module):
    def __init__(self):
        super(main_class, self).__init__()        
        self.paramA = nn.Conv2d(...)
        self.paramB = nn.Conv2d(...)
        self.in_class = inside()

    def forward(self, x):
        ...
        return x

device = #Define what GPU device to use or CPU

object = main_class()
object = object.to(device)
 

Предположим, что в этом коде device используется GPU2. Тогда я знаю, что параметры self.paramA и self.paramB определенно были загружены в GPU2, а не на CPU или любом другом графическом процессоре. Но что можно сказать о self.weight_h и self.weight_v ? Гарантированно ли они будут на GPU2 или мне нужно явно указать это для параметров inside класса?

Я использую PyTorch 1.8.1, но, возможно, предложу метод, который является довольно общим и который будет верен для любой версии PyTorch>=1.0

Я думаю, когда вы сказали этот код — термин можно уточнить немного подробнее. Есть две вещи, которые можно поместить в графический процессор. Одна из проблем касается данных. Вы можете хранить свои данные в графическом процессоре и тому подобном.

Есть еще одна часть, модель может быть перенесена на графический процессор. В этом случае, когда вы это сделаете final_model.to(...) , все модули внутри него как часть окончательной модели будут перенесены на графический процессор.

Я различил эти две вещи, потому что иногда легко перепутать эти две вещи.

Таким образом, окончательный ответ: да, они гарантированно будут на GPU. (Те внутренние веса модели, которые являются частью большой модели).