Метод cpu

Метод cpu класса Module перемещает все параметры и буферы модели на центральный процессор (CPU). Метод вызывается непосредственно на экземпляре модуля и изменяет его состояние, возвращая сам модуль для удобства цепочечных вызовов. Это полезно при переключении между устройствами, например, после обучения на GPU для выполнения инференса на CPU или для сохранения модели.

Синтаксис

module.cpu()

Метод не принимает никаких параметров и возвращает сам модуль с перемещёнными на CPU параметрами.

Пример

Давайте создадим простую линейную модель и переместим её на CPU:

import torch
import torch.nn as nn

model = nn.Linear(5, 3)
model = model.cpu()
print(next(model.parameters()).device)

Результат выполнения кода:

cpu

Пример

Рассмотрим пример перемещения модели с GPU на CPU. Для этого сначала создадим тензор и модель на GPU, а затем применим метод cpu:

import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Linear(5, 3)
model = model.cuda()
print("Device after cuda:", next(model.parameters()).device)

model = model.cpu()
print("Device after cpu:", next(model.parameters()).device)

Результат выполнения кода:

Device after cuda: cuda:0
Device after cpu: cpu

Пример

Метод cpu также перемещает буферы модели, например, тензоры, зарегистрированные через register_buffer:

import torch
import torch.nn as nn

class MyModule(nn.Module):
	def __init__(self):
		super().__init__()
		self.linear = nn.Linear(2, 2)
		self.register_buffer('buffer', torch.ones(2, 2))

model = MyModule()
model = model.cuda()
print("Buffer device:", model.buffer.device)

model = model.cpu()
print("Buffer device after cpu:", model.buffer.device)

Результат выполнения кода:

Buffer device: cuda:0
Buffer device after cpu: cpu

Смотрите также

метод to,
который перемещает модуль на указанное устройство или меняет тип данных
метод cuda,
который перемещает модуль на графический процессор (GPU)
метод mps,
который перемещает модуль на процессор Apple Metal (MPS)
метод state_dict,
который возвращает словарь с параметрами модуля для сохранения