小白学Pytorch系列–Torch.optim API Base class(1)
torch.optim是一个实现各种优化算法的包。大多数常用的方法都已得到支持，而且接口足够通用，因此将来还可以轻松集成更复杂的方法。
如何使用优化器
使用手torch.optim您必须构造一个优化器对象，该对象将保存当前状态，并将根据计算出的梯度更新参数。
构造它
要构造一个优化器，你必须给它一个包含参数(所有应该是变量)的可迭代对象来优化。然后，您可以指定特定于优化器的选项，如学习率、权重衰减等。
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr=0.0001)
每个参数选项
优化器还支持指定每个参数的选项。要做到这一点，不是传递一个变量的迭代对象，而是传递一个dict的迭代对象。它们每个都将定义一个单独的参数组，并且应该包含一个params键，包含属于它的参数列表。其他键应该与优化器接受的关键字参数匹配，并将用作该组的优化选项。
注意： 您仍然可以将选项作为关键字参数传递。在没有覆盖它们的组中，它们将作为默认值使用。当您只想改变一个选项，同时在参数组之间保持所有其他选项一致时，这很有用。
例如，当想要指定每层的学习率时，这是非常有用的
optim.SGD([
{'params': model.base.parameters()},
{'params': model.classifier.parameters(), 'lr': 1e-3}
], lr=1e-2, momentum=0.9)
这意味着这个model.base参数将使用默认的学习率1e-2，model.classifier’参数将使用1e-3的学习率，所有参数将使用0.9的动量。
进行优化步骤
所有优化器都实现了一个step()方法，用于更新参数。它有两种用法optimizer.step()
这是大多数优化器支持的简化版本。该函数可以在梯度计算完成后调用，例如backward()。
例如：
for input, target in dataset:
optimizer.zero_grad()
output = model(input)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()
optimizer.step(closure)
一些优化算法(如共轭梯度和LBFGS)需要多次重新计算函数，所以你必须传入一个闭包，允许它们重新计算你的模型。闭包应该清除梯度，计算损失并返回。
for input, target in dataset:
def closure():
optimizer.zero_grad()
output = model(input)
loss = loss_fn(output, target)
loss.backward()
return loss
optimizer.step(closure)
Base class
这部分参考了：https://zhuanlan.zhihu.com/p/87209990
PyTorch 的优化器基本都继承于 “class Optimizer”，这是所有 optimizer 的 base class。
下面是Optimizer的结构
class Optimizer(object):
def __init__(self, params, defaults):
self.defaults = defaults
self._hook_for_profile()
if isinstance(params, torch.Tensor):
raise TypeError("params argument given to the optimizer should be "
"an iterable of Tensors or dicts, but got " +
torch.typename(params))
self.state = defaultdict(dict)
self.param_groups = []
param_groups = list(params)
if len(param_groups) == 0:
raise ValueError("optimizer got an empty parameter list")
if not isinstance(param_groups[0], dict):
param_groups = [{'params': param_groups}]
for param_group in param_groups:
self.add_param_group(param_group)
def state_dict(self):
...
def load_state_dict(self, state_dict):
...
def cast(param, value):
...
def zero_grad(self, set_to_none: bool = False):
...
def step(self, closure):
...
def add_param_group(self, param_group):
...
init 函数初始化
params和defaults是两个重要的参数，defaults定义了全局优化默认值，params定义了模型参数和局部优化默认值。
add_param_group
defaultdict的作用在于当字典里的 key 被查找但不存在时，返回的不是keyError而是一个默认值，此处defaultdict(dict)`返回的默认值会是个空字典。最后一行调用的self.add_param_group(param_group)，其中param_group是个字典，Key 就是params，Value 就是param_groups = list(params)。
def add_param_group(self, param_group):
params = param_group['params']
if isinstance(params, torch.Tensor):
param_group['params'] = [params]
elif isinstance(params, set):
raise TypeError('optimizer')
else:
param_group['params'] = list(params)
for param in param_group['params']:
if not isinstance(param, torch.Tensor):
raise TypeError("optimizer " + torch.typename(param))
if not param.is_leaf:
raise ValueError("can't optimize a non-leaf Tensor")
for name, default in self.defaults.items():
if default is required and name not in param_group:
raise ValueError("parameter group didn't specify a value of required optimization parameter " +
name)
else:
param_group.setdefault(name, default) # 给参数设置默认参数
params = param_group['params']
if len(params) != len(set(params)):
warnings.warn("optimizer contains ", stacklevel=3)
param_set = set()
for group in self.param_groups:
param_set.update(set(group['params']))
if not param_set.isdisjoint(set(param_group['params'])): # 判断两个集合是否包含相同的元素
raise ValueError("some parameters appear in more than one parameter group")
self.param_groups.append(param_group)
zero_grad
就是将所有参数的梯度置为零p.grad.zero_()。detach_()的作用是Detaches the Tensor from the graph that created it, making it a leaf. self.param_groups是列表，其中的元素是字典。
def zero_grad(self):
r"""Clears the gradients of all optimized :class:`torch.Tensor` s."""
for group in self.param_groups:
for p in group['params']:
if p.grad is not None:
p.grad.detach_()
p.grad.zero_()
step
更新参数作用， 在父类 Optimizer 的 step 函数中只有一行代码raise NotImplementedError。网络模型参数和优化器的参数都保存在列表 self.param_groups 的元素中，该元素以字典形式存储和访问具体的网络模型参数和优化器的参数。所以，可以通过两层循环访问网络模型的每一个参数 p 。获取到梯度d_p = p.grad.data之后，根据优化器参数设置是否使用 momentum或者nesterov再对参数进行调整。最后一行 p.data.add_(-group['lr'], d_p)的作用是对参数进行更新。state用于保存本次更新是优化器第几轮迭代更新参数。
下面以SGD优化器为例
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
params_with_grad = []
d_p_list = []
momentum_buffer_list = []
weight_decay = group['weight_decay']
momentum = group['momentum']
dampening = group['dampening']
nesterov = group['nesterov']
maximize = group['maximize']
lr = group['lr']
for p in group['params']:
if p.grad is not None:
params_with_grad.append(p)
d_p_list.append(p.grad)
state = self.state[p]
if 'momentum_buffer' not in state:
momentum_buffer_list.append(None)
else:
momentum_buffer_list.append(state['momentum_buffer'])
F.sgd(params_with_grad,
d_p_list,
momentum_buffer_list,
weight_decay=weight_decay,
momentum=momentum,
lr=lr,
dampening=dampening,
nesterov=nesterov,
maximize=maximize,)
# update momentum_buffers in state
for p, momentum_buffer in zip(params_with_grad, momentum_buffer_list):
state = self.state[p] ## 保存
state['momentum_buffer'] = momentum_buffer
return loss
F.sgd
def sgd(params: List[Tensor],
d_p_list: List[Tensor],
momentum_buffer_list: List[Optional[Tensor]],
*,
weight_decay: float,
momentum: float,
lr: float,
dampening: float,
nesterov: bool,
maximize: bool):
for i, param in enumerate(params):
d_p = d_p_list[i]
if weight_decay != 0:
d_p = d_p.add(param, alpha=weight_decay)
if momentum != 0:
buf = momentum_buffer_list[i]
if buf is None:
buf = torch.clone(d_p).detach()
momentum_buffer_list[i] = buf
else:
buf.mul_(momentum).add_(d_p, alpha=1 - dampening)
if nesterov:
d_p = d_p.add(buf, alpha=momentum)
else:
d_p = buf
alpha = lr if maximize else -lr
param.add_(d_p, alpha=alpha)
SGD上引入了一个Momentum（又叫Heavy Ball）的改进。
load_state_dict
加载优化器状态。文章来源：https://www.toymoban.com/news/detail-411952.html
def load_state_dict(self, state_dict):
# deepcopy, to be consistent with module API
state_dict = deepcopy(state_dict)
# Validate the state_dict
groups = self.param_groups
saved_groups = state_dict['param_groups']
if len(groups) != len(saved_groups):
raise ValueError("loaded state dict has a different number of "
"parameter groups")
param_lens = (len(g['params']) for g in groups)
saved_lens = (len(g['params']) for g in saved_groups)
if any(p_len != s_len for p_len, s_len in zip(param_lens, saved_lens)):
raise ValueError("loaded state dict contains a parameter group "
"that doesn't match the size of optimizer's group")
# Update the state
id_map = {old_id: p for old_id, p in
zip(chain.from_iterable((g['params'] for g in saved_groups)),
chain.from_iterable((g['params'] for g in groups)))}
def cast(param, value):
r"""Make a deep copy of value, casting all tensors to device of param."""
if isinstance(value, torch.Tensor):
# Floating-point types are a bit special here. They are the only ones
# that are assumed to always match the type of params.
if param.is_floating_point():
value = value.to(param.dtype)
value = value.to(param.device)
return value
elif isinstance(value, dict):
return {k: cast(param, v) for k, v in value.items()}
elif isinstance(value, container_abcs.Iterable):
return type(value)(cast(param, v) for v in value)
else:
return value
# Copy state assigned to params (and cast tensors to appropriate types).
# State that is not assigned to params is copied as is (needed for
# backward compatibility).
state = defaultdict(dict)
for k, v in state_dict['state'].items():
if k in id_map:
param = id_map[k]
state[param] = cast(param, v)
else:
state[k] = v
state_dict
以字典的形式返回优化器的状态。文章来源地址https://www.toymoban.com/news/detail-411952.html
def state_dict(self):
# Save order indices instead of Tensors
param_mappings = {}
start_index = 0
def pack_group(group):
nonlocal start_index
packed = {k: v for k, v in group.items() if k != 'params'}
param_mappings.update({id(p): i for i, p in enumerate(group['params'], start_index)
if id(p) not in param_mappings})
packed['params'] = [param_mappings[id(p)] for p in group['params']]
start_index += len(packed['params'])
return packed
param_groups = [pack_group(g) for g in self.param_groups]
# Remap state to use order indices as keys
packed_state = {(param_mappings[id(k)] if isinstance(k, torch.Tensor) else k): v
for k, v in self.state.items()}
return {
'state': packed_state,
'param_groups': param_groups,
}
到了这里，关于小白学Pytorch系列--Torch.optim API Base class(1)的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！
