PyTorch 문서) PyTorch Recipes - Extension points in nn.Module for load_state_dict and tensor subclasses

nn.Module의 확장 기능들

• 원글 : Extension points in nn.Module for load_state_dict and tensor subclasses
• 이 글은 nn.Module에 통합된 두 가지 확장 기능들과 새로운 유틸리티 함수인 torch.utils.swap_tensors를 소개한다.
  • nn.Module에 통합된 기능은 (1) nn.Module.to()와 관련 메소드, (2) nn.Module.load_state_dict()다.
• 단, 이 글에 나온 기능은 PyTorch 2.3.0 이후의 버전에서만 사용 가능하다.

---

torch.utils.swap_tensors

• torch.utils.swap_tensors (이하 swap_tensors)는 두 개의 파이썬 텐서를 입력 받아, 두 값을 바꾸는(스왑 swap)하는 기능을 제공하는 유틸리티 함수다.

import torch
import torch.nn as nn
t1 = torch.arange(2)
t2 = torch.arange(3)
print(f"Before swapping, t1: {t1}, t2: {t2}")
torch.utils.swap_tensorw(t1, t2)
print(f"After swapping, t1: {t1}, t2: {t2}")

출력결과

Before swapping, t1: tensor([0, 1]), t2: tensor([0, 1, 2])
After swapping, t1: tensor([0, 1, 2]), t2: tensor([0, 1])