[12-01] 코드

1. ResNet 코드

import os
import torch
from torchvision import models, transforms
from torch.utils.data import DataLoader
from torch.optim import AdamW
from torchvision.models import resnet50, ResNet50_Weights
import torch.nn as nn
from datasets import load_dataset
from PIL import Image

# GPU 설정
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Device:", device)

# 데이터셋 로드
train_dataset = load_dataset("xodhks/EmoSet118K_", split="train")
test_dataset = load_dataset("xodhks/ugrp-survey-test", split="train")

# 테스트 데이터셋의 유효 라벨 목록
test_valid_label_indices = [0, 1, 2, 3, 4, 5]  # Children_Sketch에 존재하는 라벨 인덱스만 포함 

# ResNet 모델 로드 (예: ResNet-50)
weights = ResNet50_Weights.IMAGENET1K_V1
model = resnet50(weights=weights)
model = model.to(device)

# 모델 저장을 위한 디렉토리 생성
os.makedirs("top_models", exist_ok=True)
top_models = []

# 데이터셋 전처리 설정
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 데이터셋 전처리 함수
def preprocess_data(batch):
    # 이미지가 이미 PIL 객체인 경우
    if isinstance(batch['image'], list):
        batch['image'] = [img.convert("RGB") if isinstance(img, Image.Image) else Image.open(img).convert("RGB") for img in batch['image']]
    elif isinstance(batch['image'], Image.Image):
        batch['image'] = batch['image'].convert("RGB")
    else:
        batch['image'] = Image.open(batch['image']).convert("RGB")
    
    # transform 적용
    batch['image'] = [transform(img) for img in batch['image']]
    return batch

train_dataset = train_dataset.with_transform(preprocess_data)
test_dataset = test_dataset.with_transform(preprocess_data)

# DataLoader 설정
def collate_fn(batch):
    images = [item['image'] for item in batch]
    labels = [item['label'] for item in batch]

    images = torch.stack(images, dim=0)
    labels = torch.tensor(labels, dtype=torch.long)
    
    return images, labels

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn, num_workers=2)


criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-2)  

# 평가 함수
def evaluate(model, data_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in data_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, preds = torch.max(outputs, 1)

            for pred, label in zip(preds, labels):
                if pred.item() in test_valid_label_indices:
                    if pred.item() == label.item():
                        correct += 1
                total += 1

    accuracy = 100 * correct / total
    return accuracy

# 모델 저장 함수
def save_top_models(epoch, accuracy, model, top_models):
    model_filename = f"model_epoch_{epoch + 1}_accuracy_{accuracy:.2f}_EmoSet.pth"
    model_path = os.path.join("top_models", model_filename)
    top_models.append((accuracy, model_path))
    top_models = sorted(top_models, key=lambda x: x[0], reverse=True)[:10]
    torch.save(model.state_dict(), model_path)
    print("\nTop 10 Models (by accuracy):")
    for i, (acc, path) in enumerate(top_models, 1):
        print(f"Rank {i}: Accuracy = {acc:.2f}%, Model Path = {path}")
    return top_models

# 학습 루프
num_epochs = 20
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}")
    test_accuracy = evaluate(model, test_loader)
    print(f"Test Accuracy after Epoch {epoch+1}: {test_accuracy:.2f}%")
    top_models = save_top_models(epoch, test_accuracy, model, top_models)

print("Finished Training")

수정

import os
import torch
from torchvision import models, transforms
from torch.utils.data import DataLoader, WeightedRandomSampler
from torch.optim import Adam
from torchvision.models import resnet50, ResNet50_Weights
import torch.nn as nn
from datasets import load_dataset
from PIL import Image
from torch.optim.lr_scheduler import StepLR

# GPU 설정
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Device:", device)

# 데이터셋 로드
train_dataset = load_dataset("crawling-emotions-in-google-train", split="train")
test_dataset = load_dataset("JANGJIWON/UGRP_sketchset_textbook", split="train")

# 테스트 데이터셋의 유효 라벨 목록
test_valid_label_indices = [0, 1, 2, 3, 4, 5] 

# ResNet50 모델 로드 및 수정
model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
model.fc = nn.Linear(model.fc.in_features, 6)  # 클래스 수 6개로 수정
model = model.to(device)

# 모델 저장을 위한 디렉토리 생성
os.makedirs("top_models", exist_ok=True)
top_models = []

# 데이터셋 전처리 설정
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 데이터셋 전처리 함수
def preprocess_data(example):
    image = example['image']
    if isinstance(image, Image.Image):
        image = image.convert("RGB")
    return {'image': transform(image)}

# 데이터셋 전처리 적용
train_dataset = train_dataset.map(preprocess_data)
test_dataset = test_dataset.map(preprocess_data)

# WeightedRandomSampler 설정
train_labels = [example['label'] for example in train_dataset]
class_counts = torch.bincount(torch.tensor(train_labels))
class_weights = 1. / class_counts.float()
sample_weights = [class_weights[label] for label in train_labels]
sampler = WeightedRandomSampler(weights=sample_weights, num_samples=len(sample_weights), replacement=True)

# DataLoader 설정
def collate_fn(batch):
    images = [item['image'] for item in batch]
    labels = [item['label'] for item in batch]
    
    # 텐서로 변환
    images = torch.stack(images, dim=0)
    labels = torch.tensor(labels, dtype=torch.long)
    
    return images, labels

train_loader = DataLoader(train_dataset, batch_size=32, sampler=sampler, collate_fn=collate_fn, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn, num_workers=2)

# 손실 함수 및 옵티마이저 정의
criterion = nn.CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=1e-4)
scheduler = StepLR(optimizer, step_size=5, gamma=0.1)  # 학습률 감소 설정

# 평가 함수
def evaluate(model, data_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in data_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, preds = torch.max(outputs, 1)

            for pred, label in zip(preds, labels):
                if pred.item() in test_valid_label_indices:
                    if pred.item() == label.item():
                        correct += 1
                total += 1

    accuracy = 100 * correct / total
    return accuracy

# 모델 저장 함수
def save_top_models(epoch, accuracy, model, top_models):
    model_filename = f"model_epoch_{epoch + 1}_accuracy_{accuracy:.2f}.pth"
    model_path = os.path.join("top_models", model_filename)
    top_models.append((accuracy, model_path))
    top_models = sorted(top_models, key=lambda x: x[0], reverse=True)[:10]
    torch.save(model.state_dict(), model_path)
    print("\nTop 10 Models (by accuracy):")
    for i, (acc, path) in enumerate(top_models, 1):
        print(f"Rank {i}: Accuracy = {acc:.2f}%, Model Path = {path}")
    return top_models

# 학습 루프
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    scheduler.step()  # 학습률 업데이트

    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}")
    test_accuracy = evaluate(model, test_loader)
    print(f"Test Accuracy after Epoch {epoch+1}: {test_accuracy:.2f}%")
    top_models = save_top_models(epoch, test_accuracy, model, top_models)

print("Finished Training")

Crawling-train-test - 권태완

import os
import torch
from torchvision import transforms
from torch.utils.data import DataLoader
from torch.optim import AdamW
from torchvision.models import resnet50, ResNet50_Weights
import torch.nn as nn
from datasets import load_dataset
from torch.optim.lr_scheduler import StepLR

# GPU 설정
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Device:", device)

# 사용자로부터 모델 저장 경로 입력받기
save_dir = input("Enter the directory to save models (default: 'models'): ").strip()
if not save_dir:
    save_dir = "models"  # 기본 디렉토리 설정
os.makedirs(save_dir, exist_ok=True)

# 데이터셋 로드
train_dataset = load_dataset("xodhks/crawling-emotions-in-google-train", split="train")
test_dataset = load_dataset("xodhks/ugrp-survey-test", split="train")

# ResNet50 모델 로드 및 수정
model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
model.fc = nn.Linear(model.fc.in_features, 6)  # 클래스 수 6개로 수정
model = model.to(device)

# 데이터 전처리 설정
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 데이터셋 전처리 함수
def preprocess_data(example):
    image = example['image']
    if isinstance(image, torch.Tensor):
        return {'image': transform(image)}
    return {'image': transform(image.convert("RGB")), 'label': example['label']}

# 데이터셋 전처리 적용
train_dataset = train_dataset.map(preprocess_data)
test_dataset = test_dataset.map(preprocess_data)

# DataLoader 설정
def collate_fn(batch):
    images = torch.stack([item['image'] for item in batch])
    labels = torch.tensor([item['label'] for item in batch], dtype=torch.long)
    return images, labels

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn)

# 손실 함수 및 옵티마이저 정의
criterion = nn.CrossEntropyLoss()
optimizer = AdamW(model.parameters(), lr=1e-4)
scheduler = StepLR(optimizer, step_size=5, gamma=0.1)

# 평가 함수
def evaluate(model, data_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in data_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, preds = torch.max(outputs, 1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)
    accuracy = 100 * correct / total
    return accuracy

# 학습 루프
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    scheduler.step()

    train_accuracy = evaluate(model, train_loader)
    test_accuracy = evaluate(model, test_loader)

    print(f"Epoch [{epoch+1}/{num_epochs}]")
    print(f"  Loss: {running_loss / len(train_loader):.4f}")
    print(f"  Train Accuracy: {train_accuracy:.2f}%")
    print(f"  Test Accuracy: {test_accuracy:.2f}%")

    # 모델 저장
    model_path = os.path.join(save_dir, f"resnet50_epoch_{epoch+1}_acc_{test_accuracy:.2f}.pth")
    torch.save(model.state_dict(), model_path)

print("Training Complete")