PyTorch - Optimizer (1)

steadycode·2022년 11월 20일
0

motivation

PyTorch 에서 제공하는 optimizer 는 여러가지가 존재한다. 그러나 모두 step() 함수를 사용하여 최적화를 진행하는 것은 동일하다. DNN 학습과정을 간단히 하면 다음과 같다.

net = DNN()
input = input_data(); target = target_data()
optimizer = opt(net.parameters())
# training
for i in range(0, epoch):
	x = input
    output = net(x)
    loss = loss(output, target)
    loss.backward()
    optimizer.step()

그렇다면 여기에서 optimizer.step 함수가 net이 가지고 있는 파라미터에 어떻게 접근하여 최적화를 진행을 할까? 본 글은 이 점에 대해 다룬다.


Adam Optimizer

대표적인 Adam 의 경우 다음과 같은 step() 함수를 가지고 있다.

	@torch.no_grad()
    def step(self, closure=None):
        """Performs a single optimization step.

        Args:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        self._cuda_graph_capture_health_check()

        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            params_with_grad = []
            grads = []
            exp_avgs = []
            exp_avg_sqs = []
            max_exp_avg_sqs = []
            state_steps = []
            beta1, beta2 = group['betas']

            for p in group['params']:
                if p.grad is not None:
                    params_with_grad.append(p)
                    if p.grad.is_sparse:
                        raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
                    grads.append(p.grad)

                    state = self.state[p]
                    # Lazy state initialization
                    if len(state) == 0:
                        state['step'] = torch.zeros((1,), dtype=torch.float, device=p.device) \
                            if self.defaults['capturable'] else torch.tensor(0.)
                        # Exponential moving average of gradient values
                        state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                        # Exponential moving average of squared gradient values
                        state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
                        if group['amsgrad']:
                            # Maintains max of all exp. moving avg. of sq. grad. values
                            state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                    exp_avgs.append(state['exp_avg'])
                    exp_avg_sqs.append(state['exp_avg_sq'])

                    if group['amsgrad']:
                        max_exp_avg_sqs.append(state['max_exp_avg_sq'])

                    state_steps.append(state['step'])

            adam(params_with_grad,
                 grads,
                 exp_avgs,
                 exp_avg_sqs,
                 max_exp_avg_sqs,
                 state_steps,
                 amsgrad=group['amsgrad'],
                 beta1=beta1,
                 beta2=beta2,
                 lr=group['lr'],
                 weight_decay=group['weight_decay'],
                 eps=group['eps'],
                 maximize=group['maximize'],
                 foreach=group['foreach'],
                 capturable=group['capturable'])

        return loss

위 과정을 글로 간단히 요약하면 다음과 같다.

  1. 다음의 값을 list 로 생성한다.
    • params_with_grad
    • grads
    • exp_avgs
    • exp_avg_sqs
    • max_exp_avg_sqs
    • state_steps
    • beta1, beta2
  2. self.param_groups 에서 "param" 값을 검색한다.
  3. p.grad 가 None 값이 아니면 param_with_grad grads 에 parameter 를 저장한다.
  4. (1) 값을 모아 adam() 함수를 호출한다.
  5. loss 값을 리턴한다.

optimizer track the model parameters

그렇다면 이러한 optimizer 가 모델의 파라미터 값들을 어떻게 추론할 수 있을까? 서로 다른 클래스 객체가 정의되었음에도, optimizer는 model 의 파라미터를 간단히 self.param_groups로 추론한다. 처음에 어떻게 초기화가 되었을지 궁금하다.

# optimizer init
optimizer = opt(net.parameters())

이렇게 opt 함수를 부르면 다음의 형식으로 optimizer 객체를 초기화시킨다.
1. net.parameters() 값을 params 형태로 받음.
2. 위 값을 self.add_param_group 의 함수로 self.param_groups 에 저장함.

def __init__(self, params, defaults):
        torch._C._log_api_usage_once("python.optimizer")
        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)

        # Allows _cuda_graph_capture_health_check to rig a poor man's TORCH_WARN_ONCE in python,
        # which I don't think exists
        # https://github.com/pytorch/pytorch/issues/72948
        self._warned_capturable_if_run_uncaptured = True

python은 일반적으로 return 을 했을 경우, 값만 리턴하지 주소값을 리턴하지 않는다. 그렇다면 이것을 어떻게 net의 파라미터와 연동될 수 있게 할까

parameters()

net.parameters() 를 출력해보면 다음과 같이 나온다.

<generator object Module.parameters at ...>

위와 같이 parameter 값 자체를 리턴하지 않고 generator 라는 객체를 리턴한다. 함수를 더 자세히 분석해보면 다음과 같다.

    def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
        r"""Returns an iterator over module parameters.

        This is typically passed to an optimizer.

        Args:
            recurse (bool): if True, then yields parameters of this module
                and all submodules. Otherwise, yields only parameters that
                are direct members of this module.

        Yields:
            Parameter: module parameter

        Example::

            >>> for param in model.parameters():
            >>>     print(type(param), param.size())
            <class 'torch.Tensor'> (20L,)
            <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

        """
        for name, param in self.named_parameters(recurse=recurse):
            yield param

일반적인 함수처럼 return 명령어를 사용하는 것이 아닌, yield 를 사용하는 것을 확인할 수 있다.

yield

그렇다면 yield 가 포인터를 전달하는 것일까? 이를 탐구하기 위해 다음과 같은 테스트 코드를 짰다.

test 1

class TMP():
	def __init__(self):
    	self.init = [1,2,3,4,5,6]
    def generator(self):
    	for a in self.init:
        	yield a
if __name__ == "__main__":
	tmp = TMP()
	test = tmp.generator()
    print(test)
    print("before append")
    for a in test:
    	print(a)
    tmp.init.append(7)
    print(test)
    print("after append")
    for a in test:
    	print(a)
    print(tmp.init)

output

<generator object Tmp.generater at 0x7f415c949c10>
before append
1
2
3
4
5
6
<generator object Tmp.generater at 0x7f415c949c10>
after append
[1, 2, 3, 4, 5, 6, 7]

예상과 다른 결과를 얻을 수 있었다. pointer를 전달한다면 after append 이후 7이 포함된 리스트를 출력해야하는데 그러지 않았다. 반면 generator object의 주소값은 동일한 것을 확인했다.

test 2

반면, pytorch 의 optimizer 의 경우 처음 input으로 받은 generator() 로 parameter의 변경값을 성공적으로 추적했다. 다음과 같은 코드 블럭을 보면, self.param_groups 변수를 활용하는 것을 알 수 있다.

for group in self.param_groups:
...

그렇다면, 위 self.param_groups 변수를 어떻게 정의했기에 성공적으로 추적할 수 있었을까? generatorself.param_groups 에 저장하는 add_param_group 함수를 살펴보자

    def add_param_group(self, param_group):
        r"""Add a param group to the :class:`Optimizer` s `param_groups`.

        This can be useful when fine tuning a pre-trained network as frozen layers can be made
        trainable and added to the :class:`Optimizer` as training progresses.

        Args:
            param_group (dict): Specifies what Tensors should be optimized along with group
                specific optimization options.
        """
        assert isinstance(param_group, dict), "param group must be a dict"

        params = param_group['params']
        if isinstance(params, torch.Tensor):
            param_group['params'] = [params]
        elif isinstance(params, set):
            raise TypeError('optimizer parameters need to be organized in ordered collections, but '
                            'the ordering of tensors in sets will change between runs. Please use a list instead.')
        else:
            param_group['params'] = list(params)

        for param in param_group['params']:
            if not isinstance(param, torch.Tensor):
                raise TypeError("optimizer can only optimize Tensors, "
                                "but one of the params is " + 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 a parameter group with duplicate parameters; "
                          "in future, this will cause an error; "
                          "see github.com/pytorch/pytorch/issues/40967 for more information", 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)

특징을 살펴보면 다음과 같다.
1. 전달받는 데이터를 dict() 형태로 강제함.
2. 전달받은 dict() 데이터에 다시 list 형태로 element 저장
3. 이를 append()self.param_groups 에 저장
4. 클래스 형태로 데이터를 받음

위 과정을 다시 코드로 짜보았다. 추가적으로 Tensor() 오브젝트를 정의했다.

class Tensor():
    def __init__(self, num):
        self.num = num
        self.grad = False

class TMP():
    def __init__(self):
        self.init = [Tensor(1), Tensor(2)] 
    
    def generater(self):
        for a in self.init:
            yield a
            
if __name__ == "__main__":
    tmp = TMP() # get generator
    output_list = []
    
    tmp_list = list(tmp.generater()) 
    gen_dict = {"gen": tmp_list} # dict datatype
    
    gen_list = gen_dict['gen'] # get list
    gen_dict = gen_list # restore gen_list
    
    output_list.append(gen_dict)
    
    print(output_list)
    
    print("before append")
    for o in output_list:
        for p in o:
            print(p.grad, p.num)
    
    tmp.init[1].grad = True
    for o in output_list:
        for p in o:
            print(p.grad, p.num)

결과를 보니 성공적으로 값이 변경된 것을 확인했다.

[[<__main__.Tensor object at 0x7ff993b2b040>, <__main__.Tensor object at 0x7ff993a14700>]]
before append
False 1
False 2
False 1
True 2

conclusion

위와 같은 generator 들은 값과 더불어서 object의 주소값을 리턴한다. 따라서 optimizer가 변수를 추적할 수 있는 이유는 이러한 tensor object의 위치 값을 받아 추적하기 때문에 가능한 것으로 결론 내렸다.

profile
steadycode

0개의 댓글