# 데이터 확인하기
train_labels = os.listdir("train")
print(train_labels)
print(len(train_labels))
결과
# 데이터 확인하기
valid_labels = os.listdir("validation")
print(valid_labels)
print(len(valid_labels))
결과
import shutil
# valid 데이터에서 train 데이터에 없는 것을 삭제
for val_label in valid_labels:
if val_label not in train_labels:
shutil.rmtree(os.path.join("validation", val_label))
# 삭제 후 확인
valid_labels = os.listdir("validation")
len(valid_labels)
결과
# train 데이터에는 있지만 valid 데이터에 없는 것을 확인
for train_label in train_labels:
if train_label not in valid_labels:
print(train_label)
결과
# train 데이터에는 있지만 valid 데이터에 없는 디렉토리 파일을 만들기
# 데이터의 양을 맞춰주기 위함이다.
for train_label in train_labels:
if train_label not in valid_labels:
os.makedirs(os.path.join("validation", train_label),exist_ok=True)
# 만들고 확인하기
valid_labels = os.listdir("validation")
len(valid_labels)
결과
# 필요한 모듈 import
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader
# device 확인
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
결과
data_transforms = {
"train": transforms.Compose([
transforms.Resize((224, 224)), # 사이즈를 바꿔줌, 사이즈 통일
# 회전 각도 0, 이미지 기울기 최대 10도까지 변형 적용, 이미지 80% ~ 120% 사이 임의의 크기로 확대/축소
transforms.RandomAffine(0, shear=10, scale=(0.8, 1.2)),
transforms.RandomHorizontalFlip(), # 이미지를 랜덤으로 수평으로 뒤집음
transforms.ToTensor() # Tensor형으로 변환
]),
"validation": transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor()
])
}
# 위에 함수를 사용하여 폴더에 있는 사진들 변환
image_datasets = {
"train": datasets.ImageFolder("train", data_transforms["train"]),
"validation": datasets.ImageFolder("validation", data_transforms["validation"])
}
- 구글의 연구팀이 개발한 이미지 분류, 객체 검출 등 컴퓨터 비전 작업에서 높은 성능을 보여주는 신경망 모델이다. - 신경망의 깊이, 너비, 해상도를 동시에 확장하는 방법을 통해 효율성과 성능을 극대화한 것이 특징이다. - EfficientnetB4는 EfficientNet 시리즈의 중간 크기 모델이다.
# 필요한 모듈 import
from torchvision.models import efficientnet_b4, EfficientNet_B4_Weights
from torchvision.models._api import WeightsEnum
from torch.hub import load_state_dict_from_url
def get_state_dict(self, *args, **kwargs):
kwargs.pop("check_hash")
return load_state_dict_from_url(self.url, *args, **kwargs)
WeightsEnum.get_state_dict = get_state_dict
# 모델 생성
model = efficientnet_b4(weights=EfficientNet_B4_Weights.IMAGENET1K_V1).to(device)
model
결과
for param in model.parameters():
param.requires_grad = False
model.classifier = nn.Sequential(
nn.Linear(1792, 512),
nn.ReLU(),
nn.Linear(512, 149)
).to(device)
print(model)
결과
optimizer = optim.Adam(model.classifier.parameters(), lr=0.001)
epochs = 10
for epoch in range(epochs):
for phase in ['train', 'validation']:
if phase == 'train':
model.train()
else:
model.eval()
sum_losses = 0
sum_accs = 0
for x_batch, y_batch in dataloaders[phase]:
x_batch = x_batch.to(device)
y_batch = y_batch.to(device)
y_pred = model(x_batch)
loss = nn.CrossEntropyLoss()(y_pred, y_batch)
if phase == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
sum_losses = sum_losses + loss
y_prob = nn.Softmax(1)(y_pred)
y_pred_index = torch.argmax(y_prob, axis=1)
acc = (y_batch == y_pred_index).float().sum() / len(y_batch) * 100
sum_accs = sum_accs + acc
avg_loss = sum_losses / len(dataloaders[phase])
avg_acc = sum_accs / len(dataloaders[phase])
print(f'{phase:10s}: Epoch {epoch+1:4d}/{epochs} Loss: {avg_loss:.4f} Accuracy: {avg_acc:.2f}%')
결과
# 학습된 모델 파일 저장
torch.save(model.state_dict(), "model.pth") # model.h5
