File size: 808 Bytes
5407bdc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | """Shared base class.
``torch.optim.Optimizer`` is used strictly as a base class, as the assignment
allows: it supplies parameter-group bookkeeping and the ``state`` dict, and
every update rule below is written out by hand.
"""
import torch
class BaseOptimizer(torch.optim.Optimizer):
@staticmethod
def _closure(closure):
if closure is None:
return None
with torch.enable_grad():
return closure()
def memory_bytes(self) -> int:
"""Bytes held in optimizer state -- the number that matters for the
memory-efficient category."""
total = 0
for st in self.state.values():
for v in st.values():
if torch.is_tensor(v):
total += v.numel() * v.element_size()
return total
|