torch.nn.Module.parameters()

개발하는 G0·2023년 11월 22일
0

https://easy-going-programming.tistory.com/11

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

torch.nn.Parameter 클래스는 torch.Tensor 클래스를 상속받아 만들어졌고, torch.nn.Module 클래스의 attribute로 할당하면, 자동으로 파라메터 리스트에 추가되는 것이 기존의 torch.Tensor 클래스와의 차이점이라고 한다.

그런데 신경망 구조를 구현할때 보통 아래와 같이 하위 모듈을 추가하는식으로 구현을 하기 때문에, 파라메터를 직접 추가할 일이 거의 없었다.

class MyModule(torch.nn.Module):
    def __init__(self):
    	super(MyModule, self).__init__()
        self.A = torch.nn.Linear(100, 200)
        self.B = torch.nn.Linear(200, 10)
 	
    def forward(input):
    	return self.B(self.A(input))

그렇다면 하위 모듈만 추가해도 자동으로 파라메터가 설정되는 마법(?)은 어떻게 구현된 것일까?
공식 문서에는 설명이 없었기 때문에 torch.nn.Module 클래스의 코드를 직접 열어서 읽어보았다.

    def __init__(self):
        """
        Initializes internal Module state, shared by both nn.Module and ScriptModule.
        """
        torch._C._log_api_usage_once("python.nn_module")

        self.training = True
        self._parameters = OrderedDict()
        self._buffers = OrderedDict()
        self._backward_hooks = OrderedDict()
        self._forward_hooks = OrderedDict()
        self._forward_pre_hooks = OrderedDict()
        self._state_dict_hooks = OrderedDict()
        self._load_state_dict_pre_hooks = OrderedDict()
        self._modules = OrderedDict()

위 코드는 torch.nn.Module 클래스의 생성자를 정의하는데, _parameters와 _modules 라는 딕셔너리 타입의 멤버변수들이 있는것을 볼 수 있다. 당연하지만 _paramters는 파라메터들이, _modules에는 하위 모듈들이 들어간다. (key: 이름, value: 인스턴스)

조금 더 코드를 읽어본 결과, parameters() 메소드는 먼저 모든 하위 모듈들을 탐색하고(recursive=True), 각 모듈의 _parameters에 들어있는 파라메터들을 하나씩 반환해주는 함수였다.

아래의 그림은 모듈의 한 예시인데, 이 경우 A.parameters()를 호출하면, [A.B.weight, A.B.bias, A.C.D.weight, A.C.D.bias]가 반환된다. (물론 recursive=False를 지정하면 직접 A에 속한 파라메터만 반환하는데, 아래 경우에는 아무것도 반환되지 않게된다.)

profile
초보 개발자

0개의 댓글