Xiaobai は Pytorch シリーズを学習します -- Torch.optim API 基本クラス(1)

Xiaobai は Pytorch シリーズを学習します – Torch.optim API 基本クラス (1)

torch.optimは、さまざまな最適化アルゴリズムを実装するパッケージです。一般的に使用されるメソッドのほとんどは既にサポートされており、インターフェイスは十分に汎用的であるため、将来、より複雑なメソッドを簡単に統合できます。

オプティマイザの使用方法

手で torch.optim を使用して、現在の状態を保持し、計算された勾配に基づいてパラメーターを更新するオプティマイザー オブジェクトを構築する必要があります。

それを構築する

オプティマイザーを構築するには、最適化するパラメーター (すべて変数である必要があります) を含む反復可能なオブジェクトを与える必要があります。その後、学習率、重み減衰などのオプティマイザー固有のオプションを指定できます。

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr=0.0001)

各パラメーターオプション

オプティマイザーは、各パラメーターのオプションの指定もサポートしています。これを行うには、変数の iterable を渡す代わりに、dicts の iterable を渡します。それぞれが個別のパラメーター グループを定義し、それに属するパラメーターのリストを含む params キーを含める必要があります。他のキーは、オプティマイザが受け入れるキーワード引数と一致する必要があり、このグループの最適化オプションとして使用されます。

注: オプションをキーワード引数として渡すこともできます。それらをオーバーライドしないグループでは、デフォルトとして使用されます。これは、パラメーター グループ間で他のすべてのオプションの一貫性を維持しながら、1 つのオプションのみを変更する場合に便利です。

層ごとに学習率を指定したい場合などに便利です。

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’使用し、すべてのパラメータは運動量 0.9 を使用します。1e-3

最適化手順を実行する

すべてのオプティマイザは、step()パラメータを更新するメソッドを実装しています。これには 2 つの用途があり、
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)

基本クラス

この部分の参照先: https://zhuanlan.zhihu.com/p/87209990
PyTorch のオプティマイザーは基本的に、すべてのオプティマイザーの基本クラスである「クラス オプティマイザー」から継承します。
オプティマイザの構造は次のとおりです。

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と はdefaults2 つの重要なパラメーターです。defaults はグローバル最適化のデフォルト値を定義し、params はモデル パラメーターとローカル最適化のデフォルト値を定義します。

add_param_group

defaultdictkey 被查找但不存在时，返回的不是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_()The role of 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_()

ステップ

パラメータを更新する関数については、親クラス Optimizer のstep関数raise NotImplementedError。ネットワーク モデルのパラメーターとオプティマイザーのパラメーターは listself.param_groupsのこの要素は、ネットワーク モデルの特定のパラメーターとオプティマイザーのパラメーターをディクショナリの形式で格納してアクセスします。したがって、ネットワーク モデルの各パラメーターは、2 層の反復を通じてアクセスできますp。勾配を取得したらd_p = p.grad.data、オプティマイザーのパラメーター設定が運動量を使用するかネステロフを使用するかに従って、パラメーターを調整します。最後の行p.data.add_(-group['lr'], d_p)の機能は、パラメーターを更新することです。状態は、この更新を保存するために使用されます。これは、パラメーターを更新するためのオプティマイザーの反復回数です。

例として 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)

Momentum (別名 Heavy Ball) の改善が SGD に導入されました。

load_state_dict

オプティマイザの状態をロードします。

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

オプティマイザの状態をディクショナリとして返します。

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,
      }