PyTorch

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

# 1. 디바이스 설정 (GPU 가용시 자동 사용)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 2. 간단한 신경망 정의
class SimpleNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, num_classes)
        )

    def forward(self, x):
        return self.layers(x)

# 3. 데이터 준비 (예시: 랜덤 데이터)
X = torch.randn(1000, 20)  # 1000개 샘플, 20 features
y = torch.randint(0, 3, (1000,))  # 3 classes

dataset = TensorDataset(X, y)
train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

# 4. 모델, 손실함수, 옵티마이저 설정
model = SimpleNet(input_size=20, hidden_size=64, num_classes=3).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 5. 학습 루프
for epoch in range(10):
    model.train()
    total_loss = 0
    for batch_x, batch_y in train_loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)

        optimizer.zero_grad()
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    print(f"Epoch {epoch+1}/10, Loss: {total_loss/len(train_loader):.4f}")

# 6. 추론
model.eval()
with torch.no_grad():
    sample = torch.randn(1, 20).to(device)
    prediction = model(sample)
    predicted_class = prediction.argmax(dim=1).item()
    print(f"Predicted class: {predicted_class}")

# 7. 모델 저장 및 로드
torch.save(model.state_dict(), "model.pt")
# 로드: model.load_state_dict(torch.load("model.pt"))

# 8. PyTorch 2.0 torch.compile() 최적화 (선택적)
# compiled_model = torch.compile(model)  # 20-30% 속도 향상

import pytorch_lightning as pl
from torch.utils.data import DataLoader

class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = SimpleNet(20, 64, 3)
        self.criterion = nn.CrossEntropyLoss()

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.model(x)
        loss = self.criterion(logits, y)
        self.log("train_loss", loss, prog_bar=True)
        return loss

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=0.001)

# 학습 (GPU/TPU 자동 지원, 체크포인트 자동 저장)
trainer = pl.Trainer(max_epochs=10, accelerator="auto")
trainer.fit(LitModel(), train_dataloaders=train_loader)