저장 공간, 스트라이드, 메모리 배치 이해하기

텐서는 저장 공간과 뷰 메타데이터로 이루어져

텐서는 두 부분으로 나눌 수 있어. 하나는 자료형이 정해진 1차원 연속 메모리 블록인 저장 공간이고, 다른 하나는 그 저장 공간을 해석하는 모양·스트라이드·오프셋 같은 뷰 메타데이터야. 슬라이스하거나 전치해도 저장 공간은 보통 그대로고 뷰만 바뀌어. 그래서 슬라이싱은 O(1)이고 .transpose()도 바이트를 옮기지 않지만, 결과의 연속성은 깨질 수 있어.

스트라이드가 뜻하는 것

t.stride()는 차원마다 정수 하나를 돌려줘. 각 값은 그 차원의 인덱스를 하나 늘릴 때 저장 공간에서 몇 칸 이동해야 하는지를 뜻해. 연속 배치인 (3, 4) float 텐서의 스트라이드는 (4, 1)이야. 다음 행으로 내려가려면 원소 네 칸을, 다음 열로 가려면 한 칸을 이동해.

전치는 저장 공간을 건드리지 않고 스트라이드만 맞바꿔. (3, 4) 텐서의 모양과 스트라이드가 각각 (4, 3), (1, 4)로 바뀌는 식이야. 그래서 전치한 뒤에는 .is_contiguous()가 False가 돼. 메모리에 놓인 순서대로 원소를 걷는 방식과 차원 순서대로 걷는 방식이 더 이상 일치하지 않기 때문이야.

왜 중요할까

어떤 연산이 새 메모리를 할당하고 어떤 연산이 메타데이터만 바꾸는지 예측할 수 있어.
이미지 배치를 NHWC에서 NCHW로 바꾸는 것처럼 겉보기에는 단순한 변환이 왜 시간과 메모리를 쓰는지 추론할 수 있어.
학습이 예상보다 메모리를 30% 더 쓰는 이유도 찾기 쉬워져. 원인은 대개 의도치 않은 복사 한 번이야.

Code

텐서 밑의 계층 살펴보기·python

import torch

t = torch.arange(12).reshape(3, 4)
print(t.shape)               # torch.Size([3, 4])
print(t.stride())            # (4, 1)  — row stride 4, col stride 1
print(t.is_contiguous())     # True
print(t.untyped_storage().size())  # 12 — single contiguous blob

# Slicing — view, no copy
row = t[1]
print(row.storage_offset())  # 4 — row 1 starts at element 4 in storage
print(row.data_ptr() == t.data_ptr() + 4 * 8)  # True (8 bytes per int64)

전치가 스트라이드 변경, 저장 공간은 안 변경·python

import torch

t = torch.arange(12).reshape(3, 4)
tt = t.T
print(tt.shape)              # torch.Size([4, 3])
print(tt.stride())           # (1, 4)  — strides swapped
print(tt.is_contiguous())    # False
print(tt.data_ptr() == t.data_ptr())  # True — SAME storage

# tt.view(-1) errors. tt.reshape(-1) works (will copy under the hood).
flat = tt.contiguous().view(-1)
print(flat.data_ptr() == t.data_ptr())  # False — new allocation

메모리 회계: 텐서 무게는?·python

import torch

t = torch.randn(1024, 1024)            # float32 by default

elem_bytes = t.element_size()          # 4
n_elem = t.nelement()                  # 1,048,576
total_mb = elem_bytes * n_elem / 1024 / 1024
print(f"{total_mb:.1f} MB")            # 4.0 MB

# Half precision halves it
t16 = t.half()
print(f"{t16.element_size() * t16.nelement() / 1024 / 1024:.1f} MB")  # 2.0 MB

# Two views of the same storage do NOT double-count
view = t[:512, :]
print(view.untyped_storage().size())   # still 1,048,576 elements

Exercise

값은 같지만 스트라이드는 다른 텐서 두 개를 만들어. 하나는 연속 배치로, 다른 하나는 .transpose()만 적용한 비연속 텐서로 두면 돼. data_ptr()를 비교해 저장 공간 공유 여부를 확인하고, 두 텐서에서 sum() 축소 연산의 시간을 재 봐. 텐서가 충분히 크면 연속 배치가 측정할 수 있을 만큼 더 빨라야 해.