PyTorch-실습 : 파이토치 코드 활용하기(2)

지난 포스팅에서, 파이토치에 나와있는 예제 코드를 살펴보았는데요...

예제 코드를 다시 살펴보면,

 

1. 데이터 다루기 (데이터셋 불러오기, 변형하기)

2. 모델 준비하기 (모델 구성, 옵티마이저, 비용 함수 설정)

3. 학습하기

 

이렇게 크게 3가지로 나눠볼 수 있습니다.

 

이뿐만 아니라 대부분의 딥러닝 구조들을 다루다 보면 데이터셋들만 바뀌고, 2번과 3번 과정은 크게 다르지 않다는 것을 확인할 수가 있습니다.

사실 달라진다면, 데이터 구조와 어떠한 결과값을 얻느냐에 따라 바뀌게 될 것이고, 저 기본 베이스에 더 응용을 해주는 것입니다.

 

그렇다면, 우리는 하나의 딥러닝 모델을 만들고 데이터셋만 바꾸어 준다면, 여러 가지 프로그램을 만들어 볼 수 있을 것입니다.

 

저는 이번 예제에서, 데이터 셋을 CIFAR 데이터가 아닌 우리의 데이터로 활용할 수 있게... 데이터셋만 바꾸어 보도록 하겠습니다. 지난 케라스 예제를 응용할 것인데요.. 연예인 사진들의 이미지로 변경해보도록 하겠습니다.

 

아무래도 지난번에는 데이터 숫자가 너무 없었기 때문에, 이번에는 연예인들 사진들을 많이 갖고 올 수 있게 코드를 살짝 수정하였습니다.

from urllib.request import urlopen
from bs4 import BeautifulSoup as bs
from urllib.parse import quote_plus
import os
from selenium import webdriver
from time import sleep

browser = webdriver.Chrome()
browser.implicitly_wait(2)
actors = ["강동원", "마동석", "원빈", "장동건"]
last_height = 0
for actor in actors :
    print(f'{actor} 다운로드 검색')
    naver_img_url = "https://search.naver.com/search.naver?where=image&sm=tab_jum&query="
    url = naver_img_url + quote_plus(actor)
    browser.get(url)
    i = 0
    while i < 5 and last_height < 30000:
        browser.execute_script("window.scrollTo(0, document.body.scrollHeight);")
        new_height = browser.execute_script("return document.body.scrollHeight")
        print(i, new_height)
        last_height = new_height
        i += 1
        sleep(3)
    html = browser.page_source
    soup = bs(html, "html.parser")
    img_all = soup.find_all(class_="_img")
    for i in range(len(img_all)):
        imgUrl = img_all[i]["src"]
        if imgUrl[0:5] == "https" :
            save_img_dir = f"./img/{actor}"
            with urlopen(imgUrl) as f:
                try:
                    if not os.path.exists(save_img_dir) :
                        os.makedirs(save_img_dir)
                except OSError as e:
                    print(e)
                with open(f"{save_img_dir}/{actor}_{i}.jpg","wb") as h:
                    img = f.read()
                    h.write(img)
        else:
            pass
    print(actor, "다운로드 완료")

기존 코드에서, 셀리늄(selenium)을 추가하였습니다. 네이버 이미지 자체가, 스크롤을 내려야만 계속 보여주기 때문에 셀리늄을 통해서 자동으로 마우스 스크롤을 내리는 구문을 추가하였습니다.

browser.execute_script("window.scrollTo(0, document.body.scrollHeight);")
new_height = browser.execute_script("return document.body.scrollHeight")

다음 함수들로, 마우스를 자동으로 내려줄 것이고, 내려갈 때마다, i와 last_height을 업데이트해주어서 사진을 다운로드할 수 있도록 수정하였습니다.

 

이렇게 받은 사진들을 얼굴만 추출하여 사용해도 가능하나, 우리는 데이터셋을 응용할 것이기 때문에, 이번 포스팅에서는 이러한 이미지들을 CIFAR 데이터셋을 대체해보도록 하겠습니다.

 

또한 우리는 train_data 외에도, test_data가 필요하기 때문에, 미리 분류를 시켜줍니다.

import os
import shutil

actors = ["강동원", "원빈", "장동건", "마동석"]

for actor in actors :
    save_img_dir = f"./img/{actor}"
    new_save_img_dir = f"./img_test/{actor}"
    try:
        if not os.path.exists(new_save_img_dir):
            os.makedirs(new_save_img_dir)
    except OSError as e:
        print(e)
    file_count = len([name for name in os.listdir(f"./img/{actor}/") if os.path.isfile(os.path.join(f"./img/{actor}/", name))])
    for i in range(file_count):
        try :
            if i % 5 == 1 :
                shutil.move(f"{save_img_dir}/{actor}_{i}.jpg", f"{new_save_img_dir}/{actor}_{i}.jpg")
        except :
            pass

각 폴더에, 이미지 번호가 5의 배수이면(전체 데이터에서 20%만 빼오기 위해서...) img_test 폴더로 이동해주는 코드를 구현하였습니다. 이렇게 데이터 세팅을 완료하였습니다.

 

그럼 이렇게 받은 데이터들을 변경하여 기존 코드를 수정해보도록 하겠습니다.

 

먼저, 우리가 얻은 데이터들은 사이즈가 제각각일 것이기 때문에, 사이즈를 통일시켜주도록 하겠습니다.

이때, CIFAR이미지에서는 어떠한 사이즈를 사용했는지 먼저 확인해보도록 하겠습니다.

print(images[0].size())
-> torch.Size([3, 32, 32])

이미지 한 장을 찍어 불러왔더니 32X32를 사용한다는 것을 확인하였습니다. 그럼 저도 이에 맞추어서 우리가 가져온 데이터들의 크기들을 32X32로 세팅을 해두겠습니다.

transform = transforms.Compose([transforms.Resize((24, 24)),
                           transforms.ToTensor(),
                           transforms.Normalize((0.5, 0.5, 0.5),
                                                (0.5, 0.5, 0.5))])

트랜스폼에 데이터를 가져올 때, 일괄적으로 Resize를 하였습니다.

(혹시라도 얼굴만 따로 추출하신다면, 이미지 저장하는 과정에서 resize 함수를 사용하셔도 됩니다.)

 

trainset = torchvision.datasets.ImageFolder(root= f"{os.getcwd()}/img/", transform = transform)
trainloader = torch.utils.data.DataLoader(trainset, shuffle=True, num_workers=0)
classes = trainset.classes

그럼 root 경로에 우리가 저장하였던 연예인 폴더(img)를 넣어주시면 됩니다. 그리고, 파일 개수가 서로 다르기 때문에 batch_size 조건은 빼겠습니다. (만약 두실 거면, 모든 파일 개수를 일괄적으로 맞춰 주셔야 합니다.)

testset도 똑같이 진행해보도록 하겠습니다.

testset = torchvision.datasets.ImageFolder(root= f"{os.getcwd()}/img_test/", transform = transform)
testloader = torch.utils.data.DataLoader(trainset, shuffle=True, num_workers=0)

그리고 classes도 지정을 해주어야 하는데,

classes = trainset.classes

classes = ['강동원', '마동석', '원빈', '장동건'] 를 따로 지정 안 해도 저렇게만 써줘도 됩니다!

 

그럼 나머지 부분들은 이미 에 맞추어 살짝 수정해보도록 하겠습니다.

 

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(144, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 144)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

input은 기존 제 데이터 개수에 맞게끔, 16 * 5 * 5에서, 144로 수정을 해두었습니다.

 

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(2):   # 데이터셋을 수차례 반복합니다.

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # [inputs, labels]의 목록인 data로부터 입력을 받은 후;
        inputs, labels = data

        # 변화도(Gradient) 매개변수를 0으로 만들고
        optimizer.zero_grad()

        # 순전파 + 역전파 + 최적화를 한 후
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # 통계를 출력합니다.
        running_loss += loss.item()
        if i % 10 == 1:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

loss 출력 부분도, 데이터가 많이 없으니 if i % 10 == 1: 일 때, 출력하도록 바꾸어 두었습니다.

 

저는 간단히 바꾸었지만, loss 및 옵티마이저, 모델들도 여러분들이 바꾸어서 학습을 진행하시면서 최적의 결과값을 찾으시면 좋을 것 같습니다.

 

이렇게, CIFAR 프로그램을 활용하여 연예인 이미지 분류 프로그램을 만들어 보았습니다.