이미지 데이터를 PyTorch Tensor로 보는 법

핵심 개념

1. 텐서는 이미지의 공통 언어
   LearnPyTorch 기초 파트처럼 이미지는 숫자 텐서로 표현된다. 예를 들어 RGB 이미지는 보통 [3, 224, 224]이다.
2. shape / dtype / device는 가장 먼저 확인할 디버깅 포인트
   shape가 틀리면 layer 입력이 깨지고, dtype이 틀리면 연산/메모리 문제가 나고, device가 다르면 CPU/GPU 에러가 난다.
3. PyTorch 워크플로는 Data → Model → Loss → Optimizer → Train/Inference
4. autograd는 gradient를 자동 계산한다
   loss.backward()를 호출하면 학습 가능한 파라미터에 대한 기울기가 쌓인다.
5. 학습과 추론은 분리해야 한다
   실무에서는 학습 코드와 추론 코드를 섞지 않는 습관이 중요하다. 추론 시에는 eval() + inference_mode()를 사용한다.

PyTorch 코드 예시

아래 코드는 "이미지처럼 생긴 배치"를 만들고, 아주 작은 분류 모델로 1 step 학습 + 추론을 수행한다.

import torch
import torch.nn as nn

# 재현성
torch.manual_seed(42)

# 1) 이미지 배치 만들기: [batch, channel, height, width]
images = torch.rand(8, 3, 64, 64)          # 8장의 RGB 이미지
labels = torch.randint(0, 2, (8,))         # 2클래스 라벨

device = 'cuda' if torch.cuda.is_available() else 'cpu'
images = images.to(device)
labels = labels.to(device)

print('shape:', images.shape)   # torch.Size([8, 3, 64, 64])
print('dtype:', images.dtype)   # torch.float32
print('device:', images.device) # cpu 또는 cuda

# 2) 아주 작은 이미지 분류 모델
model = nn.Sequential(
    nn.Flatten(),               # [B, 3, 64, 64] -> [B, 12288]
    nn.Linear(3 * 64 * 64, 32),
    nn.ReLU(),
    nn.Linear(32, 2)
).to(device)

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 3) 1 step 학습
model.train()
logits = model(images)                  # [8, 2]
loss = loss_fn(logits, labels)

optimizer.zero_grad()                   # 이전 gradient 초기화
loss.backward()                         # autograd가 gradient 계산
optimizer.step()                        # 파라미터 업데이트

print('train_loss:', float(loss))

# gradient 확인
first_weight_grad = model[1].weight.grad
print('grad shape:', first_weight_grad.shape)
print('grad mean:', first_weight_grad.abs().mean().item())

# 4) 추론
model.eval()
with torch.inference_mode():
    pred_logits = model(images[:2])
    pred_probs = torch.softmax(pred_logits, dim=1)
    pred_class = pred_probs.argmax(dim=1)

print('pred_probs:', pred_probs)
print('pred_class:', pred_class)

코드에서 꼭 볼 포인트

• nn.Flatten()은 이미지 텐서를 선형층에 넣기 위해 펼친다.
• CrossEntropyLoss를 쓸 때 모델 출력은 softmax 전 logits 이어도 된다.
• optimizer.zero_grad()를 빼먹으면 gradient가 누적된다.
• model.eval()과 torch.inference_mode()는 추론 시 기본 습관이다.

실무 포인트

• 채널 순서 확인: OpenCV/NumPy에서 가져오면 HWC인 경우가 많고, PyTorch 모델은 대개 CHW 또는 BCHW를 기대한다.
• dtype 기본값 정리: 이미지 학습은 대체로 float32, 라벨은 long을 많이 쓴다.
• shape 로그를 습관화: 데이터 로딩 직후와 모델 입력 직전에 shape, min/max, dtype을 확인하면 초반 디버깅 속도가 크게 빨라진다.
• 작은 배치로 먼저 통과: 실무에서는 전체 데이터 전에 4~8장으로 forward/backward가 도는지 먼저 확인하는 것이 안전하다.
• 학습/추론 분리: 추론 스크립트는 gradient 계산을 꺼서 속도와 메모리를 아껴야 한다.