torch와 tensorflow를 CNN 코드로 비교하기

Fashion MNIST 데이터 불러오기

 

torch의 경우 데이터를 불러올 때 Compose를 통해 normalize를 하고 배치 사이즈를 지정하여 iterator를 만든다.

BATCH_SIZE = 64

from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5,), std=(0.5,))
])

full_train_dataset = datasets.FashionMNIST('dataset/', train=True, download=True, transform=transform)

train_size = int(0.8 * len(full_train_dataset))
val_size = len(full_train_dataset) - train_size

train_dataset, val_dataset = random_split(full_train_dataset, [train_size, val_size])

# DataLoader 구성
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)

test_dataset = datasets.FashionMNIST('dataset/', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE)

 

tensorflow의 경우 학습할 때 배치사이즈를 정의하며 Fashoin Mnist 데이터에 대해서는 y를 원핫 인코딩을 해주어야 한다.

import tensorflow as tf
from tensorflow.keras.datasets import fashion_mnist

(X_train, y_train), (X_test, y_test) = fashion_mnist.load_data()

X_train = X_train.reshape((60000, 28, 28, 1))
X_test = X_test.reshape((10000, 28, 28, 1))

X_train = X_train.astype(float) / 255
X_test = X_test.astype(float) / 255


from tensorflow.keras.utils import to_categorical

y_train = to_categorical(y_train)
y_test = to_categorical(y_test)

from sklearn.model_selection import train_test_split

X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42)

 

 

CNN 모델 정의 및 컴파일

torch의 경우 모델을 정의할 때 일반적으로 class를 사용하며 channel first이다.

컴파일 할 때 GPU에서 쓰겠다는 것을 명시하고 옵티마이저에 모델의 파라미터를 미리 알려준다.

import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as f

class Cnn(nn.Module):
    def __init__(self):
        super(Cnn, self).__init__()
        
        self.conv1 = nn.Conv2d(1, 10, 5) # 입력 채널 수=1, 출력 채널 수=10, 커널 사이즈=5, stride=1(기본값)
        self.conv2 = nn.Conv2d(10, 20, 5)
        self.drop = nn.Dropout2d() # 출력값에 Dropout 적용
        
        # input이 320인 이유는 아래 forward의 주석 확인
        self.fc1 = nn.Linear(320, 160) # 선형결합연산을 수행하는 객체
        self.fc2 = nn.Linear(160, 80) # 히든 레이어
        self.fc3 = nn.Linear(80, 10) # 출력 레이어
    
    def forward(self, x):
        # 1x28x28 -> 10x24x24 -> 10x12x12
        x = f.relu(f.max_pool2d(self.conv1(x), 2)) # 합성곱 레이어, 풀링, 활성화 함수 순
        
        # 10x12x12 -> 20x8x8 -> 20x4x4
        x = f.relu(f.max_pool2d(self.drop(self.conv2(x)), 2)) # dropout 추가
        
        # 20x4x4 -> 1x320
        x = x.view(-1, 320) # 평탄화 작업
        x = f.relu(self.fc1(x))
        x = f.relu(self.fc2(x))
        x = f.dropout(x, training=self.training)
        x = self.fc3(x)
        return x
        
model_torch = Cnn().to(device) # 모델의 파라미터들을 지정한 장치의 메모리로 전달
optimizer = optim.RMSprop(model_torch.parameters(), lr=0.01)

 

 

한편 tensorflow의 경우 모델을 정의할 때 함수를 사용하며 channel last이다. (tensor와 비슷한 모양으로 맞추기 위해 functional api로 구현하였다.)

컴파일 할 때 loss와 optimizer, metrics를 정의한다.

from tensorflow.keras import models
from tensorflow.keras import layers

def tf_model():
    inputs = layers.Input(shape=(28, 28, 1))

    x = layers.Conv2D(filters=10, kernel_size=5, activation='relu')(inputs)
    x = layers.MaxPool2D(pool_size=(2, 2))(x)
    x = layers.Conv2D(filters=20, kernel_size=5, activation='relu')(x)
    x = layers.Dropout(0.5)(x)
    x = layers.MaxPool2D(pool_size=(2, 2))(x)

    x = layers.Flatten()(x)
    x = layers.Dense(160, activation='relu')(x)
    x = layers.Dense(80, activation='relu')(x)
    x = layers.Dropout(0.2)(x)
    outputs = layers.Dense(10, activation='softmax')(x)

    model = models.Model(inputs=inputs, outputs=outputs)

    model.compile(loss = 'categorical_crossentropy',
                  optimizer = 'rmsprop',
                  metrics = ['accuracy'])
    return model
    
model_tf = tf_model()

 

둘의 가장 큰 차이점은 tensorflow의 마지막 층을 보면 softmax가 활성화함수로 들어가있는데 반해 torch는 따로 없다. 이는 loss를 계산할 때 내부적으로 softmax를 갖고 있기 때문이다.

 

 

 

모델 학습 및 평가

개인적으로 가장 차이가 심하다고 생각하는 구간이다.

 

우선 torch의 경우

1. train, eval 모드를 정의해줘야 하고

2. 각 epoch마다 학습을 직접 실행해줘야 하고

3. loss도 직접 계산해야 하고

4. optimizer 초기화 및 역전파 계산도 직접 실행시켜줘야 하고

5. metrics에 대한 history도 직접 쌓아야 하고

6. checkpoint, earlystopping을 직접 선언해줘야 한다

는 단점들이 존재한다.

 

class EarlyStopping:
    def __init__(self, patience=5):
        self.loss = float('inf')
        self.patience = 0
        self.patience_limit = patience
        
    def step(self, loss):
        if self.loss > loss:
            self.loss = loss
            self.patience = 0
        else:
            self.patience += 1
    
    def is_stop(self):
        return self.patience >= self.patience_limit
        
def train(model, optimizer, train_iter):
    model.train()
    train_loss = 0
    # 학습 데이터를 배치 사이즈만큼 반복
    for batch_idx, (data, target) in enumerate(train_iter):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = f.cross_entropy(output, target)
        train_loss += loss.item() * len(data)
        loss.backward()
        optimizer.step()
        
    return train_loss / len(train_iter.dataset)


def evaluate(model, val_iter, len_val_data):
    model.eval() # 검증 모드
    
    loss = 0 # 테스트 데이터 오차
    correct = 0 # 정답 수
    
    with torch.no_grad(): # 평가 때는 기울기 계산 필요 없음
        for data, target in val_iter:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss += f.cross_entropy(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True) # dim=1: 차원, keepdim=True: 차원 유지
            correct += pred.eq(target.view_as(pred)).sum().item()
            
    loss /= len_val_data
    acc = correct / len_val_data * 100
    
    return loss, acc

train_history_torch = []
val_history_torch = []

best_model = None
best_val_loss = None

es = EarlyStopping(patience=10)

for epoch in range(1, EPOCHS+1):
    train_loss = train(model_torch, optimizer, train_loader)
    train_history_torch.append(train_loss)
    
    val_loss, acc = evaluate(model_torch, val_loader, len(val_dataset))
    val_history_torch.append(val_loss)
    
    print(f'[Epoch: {epoch:2d}] loss: {train_loss:.6f} | val_loss: {val_loss:.6f} - val_accuracy: {acc:.2f}%')
    
    es.step(val_loss)
    if es.is_stop():
        break
    
    if not best_val_loss or val_loss < best_val_loss:
        best_model = model_torch
        best_val_loss = val_loss
        torch.save(model_torch.state_dict(), 'bestmodel.pt')

 

위에서 한 번 설명했었지만 loss를 계산하는 함수인 `f.cross_entropy` 내부적으로 softmax를 계산하기 때문에 예측을 확률로 받고 싶다면 softmax 함수를 예측값에 취해줘야 한다.

with torch.no_grad():
    data, target = next(iter(test_loader))
    data, target = data.to(device), target.to(device)
    pred = best_model(data)[0]

print(pred)
"""
tensor([-69.9343, -78.5487, -63.8236, -75.4411, -96.9962, -26.8139, -72.7066,
        -27.9535, -47.1973, -17.0533], device='cuda:0')
"""

f.softmax(pred)
"""
tensor([1.0815e-23, 1.9626e-27, 4.8742e-21, 4.3900e-26, 1.9108e-35, 5.7674e-05,
        6.7612e-25, 1.8454e-05, 8.1020e-14, 9.9992e-01], device='cuda:0')
"""

f.softmax(pred).sum() # tensor(1.0000, device='cuda:0')

 

 

반면 tensorflow는 코드가 매우 단촐하다.

checkpoint, earlystopping이 있고 metrics와 그에 대한 history도 자동으로 남겨진다.

from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping

es = EarlyStopping(monitor='val_loss',
                   mode='min',
                   patience=10,
                   verbose=0)

mc = ModelCheckpoint('bestmodel.h5',
                     monitor='val_loss',
                     mode='min',
                     save_best_only=True,
                     verbose=0)
                     
hist_tf = model_tf.fit(X_train, y_train, epochs=EPOCHS, 
                       batch_size=BATCH_SIZE, validation_data=(X_val, y_val),
                       callbacks=[es, mc], verbose=1)

 

 

torch vs tensorflow 중 입문한다면

나는 점유율 때문에라도 torch를 추천한다. 아래 이미지는 논문들에 쓰인 딥러닝 프레임워크 점유율 그래프인데 tensorflow가 2019년 쯤부터 torch에게 따라잡히더니 최근에 와서는 점유율이 매우 낮아진 걸 볼 수 있다.

Papers With Code : Trends

tensorflow가 편한데 점유율은 왜 차이가 크게 나는 걸까?

 

검색 결과와 내 생각을 적어보자면

1. tensorflow가 편한 만큼 유연성이 떨어진다. (모델 제어가 불편하다.)

2. tensorflow 최신 버전은 windows를 지원하지 않는다. (개발자 대부분은 맥을 써서 큰 상관은 없을 수도 있지만...)

3. torch가 tensorflow에 비해 디버깅이 쉽다.

4. hugging face의 대다수 모델이 torch 기반이다.

5. torch가 tensorflow에 비해 자동 미분이 쉽다.

6. torch는 GPU에 최적화되어있는 반면 tensorflow는 TPU에 최적화되어있다.

 

특히 4번이 LLM 모델이 쏟아져나오는 요즘 시기에 큰 영향을 주지 않았을까 싶다.