Состояние оптимизатора в PyTorch
При длинном обучении полезно
запоминать не только веса модели,
но и внутренние данные спуска.
У объекта SGD метод
state_dict собирает
накопленные буферы и настройки
групп параметров. Сделаем шаг
обновления и снимем словарь:
import torch
import torch.nn as nn
torch.manual_seed(0)
model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(
model.parameters(), lr=0.1, momentum=0.9
)
x = torch.tensor([[1.0]])
y = torch.tensor([[2.0]])
pred = model(x)
loss = ((pred - y) ** 2).sum()
loss.backward()
optimizer.step()
snapshot = optimizer.state_dict()
print(len(snapshot["state"])) # выведет 2
Словарь можно записать на диск
тем же способом, что и веса модели.
На новом объекте SGD с теми
же параметрами метод load_state_dict
вернёт буферы и скорость шага:
import torch
import torch.nn as nn
torch.manual_seed(0)
model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(
model.parameters(), lr=0.1, momentum=0.9
)
x = torch.tensor([[1.0]])
y = torch.tensor([[2.0]])
pred = model(x)
loss = ((pred - y) ** 2).sum()
loss.backward()
optimizer.step()
snapshot = optimizer.state_dict()
resumed = torch.optim.SGD(
model.parameters(), lr=0.1, momentum=0.9
)
resumed.load_state_dict(snapshot)
lr = resumed.state_dict()["param_groups"][0]["lr"]
print(lr) # выведет 0.1
Создайте линейный слой 1
на 1 и объект спуска
со скоростью 0.05
и импульсом 0.8.
Выполните один шаг по паре
[[2.0]] и эталону
[[4.0]]. Снимите
словарь внутреннего состояния
спуска и выведите, сколько
записей лежит во внутреннем
состоянии.
Соберите преобразование
2 на 1, настройте
спуск с импульсом 0.9
и скоростью 0.1.
После одного шага сохраните
снимок состояния в opt.pt.
Создайте новый объект спуска
с теми же настройками для
тех же параметров, восстановите
данные из файла и выведите
скорость шага из первой группы.
Опишите модель с одним линейным слоем, выполните шаг спуска с импульсом. Перенесите внутреннее состояние на другой объект спуска того же класса и выведите, совпадает ли число записей внутреннего состояния до и после переноса.
