[PyTorch] tensor.detach()의 기능과 예시 코드

PyTorch tensor에 사용할 수 있는 detach() method는 gradient의 전파를 멈추는 역할을 한다.

 

https://pytorch.org/docs/stable/generated/torch.Tensor.detach.html

torch.Tensor.detach — PyTorch 1.13 documentation

 

다음 예시를 통해 쉽게 이해할 수 있다.

 

import torch
import torch.nn as nn

class TestModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(10, 10)
        self.layer2 = nn.Linear(10, 10)
        
    def forward(self, x):
        out1 = self.layer1(x)
        out2 = self.layer2(out1.detach())
        return out2
    
model = TestModel()

 

위의 코드를 그래프로 나타내면 다음과 같다. 

 

 

 

layer 1에서 나온 output이 detach되었기 때문에, 역전파 시 gradient가 그 이전 layer로 흘러가지 않는다.

 

각 layer의 weight gradient를 출력해보면, layer1에는 gradient가 축적되지 않은 것을 확인할 수 있다.

x = torch.randn(1, 10)
a = model(x)
a.mean().backward()
print(model.layer1.weight.grad)
print(model.layer2.weight.grad)