Метод share_memory
Метод share_memory класса Module перемещает все тензоры параметров и буферов модуля в разделяемую память. Это необходимо для эффективной работы с модулем в многопроцессных сценариях, когда несколько процессов должны иметь доступ к одним и тем же данным модели без копирования. Метод работает in-place и возвращает сам модуль, что позволяет использовать его в цепочках вызовов.
Синтаксис
module.share_memory()
Пример
Давайте создадим простую линейную модель и переместим её в разделяемую память:
import torch
import torch.nn as nn
model = nn.Linear(10, 5)
model.share_memory()
print(model.weight.is_shared())
print(model.bias.is_shared())
Результат выполнения кода:
True
True
Пример
Используем метод в цепочке вызовов сразу после создания модели:
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(20, 10),
nn.ReLU(),
nn.Linear(10, 2)
).share_memory()
for param in model.parameters():
print(param.is_shared())
Результат выполнения кода:
True
True
True
True
Пример
Проверим, что метод работает рекурсивно для всех вложенных модулей:
import torch
import torch.nn as nn
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(5, 3)
self.fc2 = nn.Linear(3, 1)
model = MyModel()
model.share_memory()
print(model.fc1.weight.is_shared())
print(model.fc2.weight.is_shared())
Результат выполнения кода:
True
True
