pytorch - custom fuction

시작하기

# 새노트 > 런타임 : GPU 지정
# 세션 새로 시작시, pytorch 새로 설치
# !pip3 install torch
# !pip3 install torchvision
import torch
import torch.nn as nn
from torch.autograd import Variable, Function

 

 

torch.autograd 모듈의 기본제공함수 확인

# torch.autograd.Function 클래스가 제공하는 함수 목록 확인
for i in Function.__dict__:
  print(i)   # Function

 

커스텀 클래스 만들기

# 사용자 정의 클래스와 메소드 만들기
class Exp(Function):
  @staticmethod
  def forward(ctx, i):
    result = i.exp()
    ctx.save_for_backward(result)
    return result

  @staticmethod
  def backward(ctx, grad_output):
    result, = ctx.saved_variables
    return grad_output * result

layer = Exp().apply

 

경사도 구하기

a = Variable(torch.Tensor([1, 2]), requires_grad=True)
output = layer(a)
print(output)