[Pytorch] Dataloader 다양하게 사용하기 (custom loader)

파이토치의 데이터로더는 아래의 형태와 같이 사용합니다.

path = 'C:/test'

dataset = torchvision.dataset.ImageFolder(root=path,
                           transform=transforms.Compose([
                               transforms.Grayscale(),
                               transforms.ToTensor(),
                           ]))

dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

이미지든 어떠한 형태의 데이터를 모두 다룰 수 있지만, 위에 사용한 ImageFolder는 이미지를 다루는 애고 다른 형식의 데이터는 조금 후에 말하겠습니다.

데이터로더는 만들어진 이유가 Classification 태스크를 다루기 위한 것이기 때문에 저 path에서 바로 데이터가 있는 것이 아니라

path -- cat -- data들

       -- dog -- data들

이런 식으로 안에 subfolder가 있고 그 안에 데이터가 있는 식으로 이루어집니다. 때문에 만약 굳이 저러한 카테고리가 필요없으신 상황이라면 그냥 대충 아무런 폴더를 만들고 그 안에 데이터를 넣으면 될 것 같습니다.

path -- 1 -- data들

위와 같은 방식으로요.

    for data in dataloader:
        img, _ = data
        img = Variable(img).to(device)

이렇게 데이터를 for 문을 통해 뽑을 수 있습니다. 

dataloader는 앞서 말한 batch size, 프로세싱 등의 파라미터를 통해 설정한 대로 우리의 폴더 내의 데이터들을 batch로 만들어서 iterative하게 반환합니다.

    for i, data in enumerate(dataloader, 0):
        print(i, end=' ')
        batch = data[0]
        print(batch.size())

이 방법 또한 데이터를 for 문을 통해 뽑을 수 있는 방법입니다. 여기서는 dataloader를 enumerate을 사용하여 돌렸는데, 뭐 방법 중 하나입니다. 

 

 

---

Custom loader 사용하기

제가 정말 애용하는 것인데, 데이터에 추가적으로 조작을 하고 싶을 때 사용할 수 있습니다.

여기서는 Datasetfolder, Custom loader 두 개념이 중요합니다.

 

위에 말한 ImageFolder 말고 Datasetfolder 가 있는데, 이는 파일 확장자를 통해서 그냥 파일들을 다 불러옵니다. 여기서 .jpg, .png 이런 식으로 하면 ImageFolder와 같은 효과를 낼 수 있겠죠.

저 같은 경우는 파일들을 추가로 조작을 해서 모두 .pt 파일로 있습니다. torch.save를 통해서 저장한 데이터들이죠.

    def custom_loader(path):
        ret = torch.load(path)
        return ret

    data_path = 'C:/200229_Prof_ds'
    dataset = dset.DatasetFolder(root=data_path,
                                 loader=custom_loader,
                                 extensions='.pt')
                                 
    dataloader = torch.utils.data.DataLoader(dataset,
                                             batch_size=4,
                                             shuffle=True,
                                             num_workers=0)

위 코드를 한 번 잘 보시면 유용하게 쓸 수 있을 겁니다. 데이터셋의 loader를 custom loader라는 함수를 만들어서 사용할 수 있습니다. 저 함수는 데이터의 path를 받아서 어떻게 하고 반환하게 되는데 데이터셋 폴더 내의 모든 파일들이 batch를 만들 때 저 함수를 거치고, concat을 통해서 batch가 만들어지고 우리가 사용하게 되는 것입니다.

def custom_loader(path):
    ret = torch.load(path)
    ret = ret.squeeze(0)
    ret = ret[:, 0:61, :, :]
    ret -= ret.min()
    ret /= ret.sum()
    # ret /= ret.max()
    return ret


def custom_image_loader(path):
    ret = Image.open(path)
    ret = to_stacked(ret)
    return ret

위와 같이 다양하게 사용할 수 있습니다. 위에 거는 인덱싱도 하고 preprocessing도 해주는 코드고 아래 거는 따로 to_stacked라는 함수를 선언해서 그 함수의 결과값을 반환하는 형식입니다.

 

저 custom loader 안에 새로 만든 다른 함수를 또 추가해도 되고 정규화를 하든 GPU로 옮기든 뭐든지 할 수 있습니다. 저 위에 extensions라는 파라미터가 확장자를 받는 곳인데 저 확장자가 아닌 파일들은 읽지 않으니 같은 폴더 내에 선택적으로 데이터를 뽑을 때에도 잘 사용하시면 될 것 같습니다.