import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torchvision import models
transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)
net = models.resnet18(pretrained=True)
num_ftrs = net.fc.in_features
net.fc = nn.Linear(num_ftrs, 10)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
net = net.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
for epoch in range(10):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 100 == 99:
print(f"[{epoch + 1}, {i + 1}] loss: {running_loss / 100:.3f}")
running_loss = 0.0
print('Finished Training')
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images.to(device), labels.to(device)
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f"Accuracy of the network on the 10000 test images: {100 * correct / total:.2f}%")
- 데이터 전처리: CIFAR-10 데이터셋에 대해 다양한 데이터 증강 및 정규화를 적용하다. 이를 통해 모델의 일반화 성능을 높이다.
- 데이터셋 로드:
torchvision
을 사용해 CIFAR-10 데이터셋을 다운로드하고, 이를 PyTorch DataLoader
에 로드하다.
- ResNet 모델 불러오기:
torchvision.models
에서 사전 훈련된 ResNet-18 모델을 불러오다. CIFAR-10에 맞게 출력층을 10개의 클래스로 수정하다.
- 모델 학습: 10 에포크 동안 SGD 옵티마이저와 교차 엔트로피 손실 함수를 사용해 모델을 학습시키다. 학습 중에는 매 100 미니 배치마다 손실을 출력하다.
- 모델 평가: 학습된 모델을 테스트 세트에 대해 평가하고, 전체 정확도를 출력하다.