PyTorch vs Tensorflow 코드 비교

목표는 파이토치로 모델을 구현하고 안드로이드용 모바일 버전 포팅 시에만 TFLite로 변환

공부하다 보니 근본적인 구조는 비슷한 것 같다.

간단한 모델 비교: PyTorch 버전

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 데이터셋 생성
x_data = torch.randn(100, 10)
y_data = torch.randint(0, 2, (100, 1)).float()
dataset = TensorDataset(x_data, y_data)
dataloader = DataLoader(dataset, batch_size=10, shuffle=True)

# 모델 정의
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.layer = nn.Linear(10, 1)
    
    def forward(self, x):
        return torch.sigmoid(self.layer(x))

model = SimpleNN()

# 손실 함수 및 최적화
criterion = nn.BCELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 학습 루프
for epoch in range(100):
    for batch_x, batch_y in dataloader:
        # 순전파
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)
        
        # 역전파 및 최적화
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch+1}/100], Loss: {loss.item():.4f}')

# 추론
with torch.no_grad():
    new_data = torch.randn(5, 10)
    predictions = model(new_data)
    print(predictions)

같은 모델의 Tensorflow 구현버전

import tensorflow as tf

# 데이터셋 생성
x_data = tf.random.normal([100, 10])
y_data = tf.random.uniform([100, 1], minval=0, maxval=2, dtype=tf.int32)

dataset = tf.data.Dataset.from_tensor_slices((x_data, y_data))
dataset = dataset.batch(10).shuffle(100)

# 모델 정의
class SimpleNN(tf.keras.Model):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.layer = tf.keras.layers.Dense(1, activation='sigmoid')
    
    def call(self, x):
        return self.layer(x)

model = SimpleNN()

# 손실 함수 및 최적화
criterion = tf.keras.losses.BinaryCrossentropy()
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)

# 학습 루프
for epoch in range(100):
    for batch_x, batch_y in dataset:
        with tf.GradientTape() as tape:
            outputs = model(batch_x)
            loss = criterion(batch_y, outputs)
        
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
    
    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch+1}/100], Loss: {loss.numpy():.4f}')

# 추론
new_data = tf.random.normal([5, 10])
predictions = model(new_data)
print(predictions)