pytorch dataloader

1. dataset 정의하고

2. loader 쓰는게 일반적인 사용법임.

 

1) Image folder 쓰는 경우

특징 :

이미지 한장씩 차곡차곡 불러오는 경우

각 폴더가 label이름임

 

예시

####tuils.py 내 내용
def imagenet_transform():
    return transforms.Compose([
        transforms.Resize((256,256)),
        transforms.RandomCrop((240,240), padding=4),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ToTensor(),     
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])


###main.py 내 내용
#같은 폴더 내 utils.py 내용 모두 가져오기
from utils import *
from tqdm import tqdm

dataset = torchvision.datasets.ImageFolder(root=f"{self.root_path}/val", transform=imagenet_ransform())
dataloader = torch.utils.data.DataLoader(
            dataset, batch_size=self.train_batch, shuffle=True, num_workers=4)

for index, (data, label) in tqdm(enumerate(loader), total=len(loader)): 
		#cuda 올리기
        data = data.cuda(device)
        #cuda 올리기 다른 방식
        data, target = data.to(device), target.to(device)

        output = model(data)

 

 

custom dataloader example1 -> 한 폴더 내 5개 이미지 불러오기

#utils.py

class CustomDataset(data.Dataset): 
    def __init__(self, args, transforms):
        self.path = args.data_path        
        self.task = os.listdir(args.data_path)
        self.task_list = [file for file in self.task if not file.endswith('jpg')]
        self.transforms = transforms


  # 총 데이터의 개수를 리턴
    def __len__(self): 
        return len(self.task_list)

    def get_image(self, path):
        # open image
        image = Image.open(path)
        image = np.array(image)
       
        
        return image  
    
    def __getitem__(self, idx): 
        
        path = os.path.join(self.path, self.task_list[idx])
        ref_list = os.listdir(path)
        ref1 = self.get_image(os.path.join(path, ref_list[0]))
        ref2 = self.get_image(os.path.join(path, ref_list[1]))
        ref3 = self.get_image(os.path.join(path, ref_list[2]))
        ref4 = self.get_image(os.path.join(path, ref_list[3]))
        ref5 = self.get_image(os.path.join(path, ref_list[4]))  
        return self.transforms(ref1),self.transforms(ref2),self.transforms(ref3),self.transforms(ref4),self.transforms(ref5)
        
  #main.py
  dataset = CustomDataset(args, transforms = imagenet_transform())
    loader = torch.utils.data.DataLoader(dataset, batch_size = 8, shuffle=False, num_workers=0)

 

custom dataloader example 2 : 데이터셋 이미지 이름순으로 소팅(natsorted) 하고 불러오기

#utils.py

class style_dataset(data.Dataset): 
    def __init__(self, args, transforms):
        self.path = args.data_path
        self.data = []
        for index, i in enumerate(natsorted(os.listdir(self.path))):
            self.data.append(i)
        self.transforms = transforms
              
    def __getitem__(self, idx):     
        img_ = os.path.join(self.path, self.data[idx])
        img = self.get_image(img_)
        return img

    def __len__(self):
        return len(self.data)
    
    def get_image(self, path):
        # open image
        image = Image.open(path)
        image = np.array(image)  
        image_tensor = self.transforms(image)

        return image_tensor
        
  #main.py
  dataset = style_dataset(args, transforms = imagenet_transform())
    loader = torch.utils.data.DataLoader(dataset, batch_size = args.batch_size, shuffle=False, num_workers=0)

 

loader 여러개 : 딕셔너리 만들고 하나하나 불러서 변수명으로 npy 저장

import torch
import torchvision
import logging
import torchvision.datasets as datasets
import timm
import numpy as np
from torchvision import transforms
import torch.nn as nn
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from _sin_reference import *
from torchvision import models


resnet50_pretrained = models.resnet50(pretrained=True)


def imagenet_valid_transform():
    return transforms.Compose([
        transforms.Resize((224,224)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
    
model_A = "resnet50_trained_on_SIN"
model_B = "resnet50_trained_on_SIN_and_IN"
model_C = "resnet50_trained_on_SIN_and_IN_then_finetuned_on_IN"
#model = load_model(model_name = model_A)

model = resnet50_pretrained

    
#data_path = "../../dataset/dongkyun/dk_dataset/afhq_dataset/experiment_afhq/for_test/stylized_dog_adained"
    
data_path = "../../dataset/dongkyun/dk_dataset/afhq_dataset/experiment_afhq/for_test/afhq_dog"
data_path_fsmr = "../../dataset/dongkyun/dk_dataset/afhq_dataset/experiment_afhq/for_test/fsmr_dog"
data_path_stylegan = "../../dataset/dongkyun/dk_dataset/afhq_dataset/experiment_afhq/for_test/stylegan_dog"

data_path_cat = "../../dataset/dongkyun/dk_dataset/afhq_dataset/experiment_afhq/for_test/afhq_cat"
data_path_dog = "../../dataset/dongkyun/dk_dataset/afhq_dataset/experiment_afhq/for_test/afhq_dog"
data_path_cutmix = "../../dataset/dongkyun/dk_dataset/afhq_dataset/experiment_afhq/for_test/cutmixed"




            
val_dataset = datasets.ImageFolder(root=f"{data_path}", transform = imagenet_valid_transform())
val_loader_default = torch.utils.data.DataLoader(val_dataset, batch_size=50, shuffle=False, num_workers=4)

val_dataset_fsmr = datasets.ImageFolder(root=f"{data_path_fsmr}", transform = imagenet_valid_transform())
val_loader_fsmr = torch.utils.data.DataLoader(val_dataset_fsmr, batch_size=50, shuffle=False, num_workers=4)

val_dataset_stylegan = datasets.ImageFolder(root=f"{data_path_stylegan}", transform = imagenet_valid_transform())
val_loader_stylegan = torch.utils.data.DataLoader(val_dataset_stylegan, batch_size=50, shuffle=False, num_workers=4)

val_dataset_cat = datasets.ImageFolder(root=f"{data_path_cat}", transform = imagenet_valid_transform())
val_loader_cat = torch.utils.data.DataLoader(val_dataset_cat, batch_size=50, shuffle=False, num_workers=4)


val_dataset_dog = datasets.ImageFolder(root=f"{data_path_dog}", transform = imagenet_valid_transform())
val_loader_dog = torch.utils.data.DataLoader(val_dataset_dog, batch_size=50, shuffle=False, num_workers=4)

val_dataset_cutmix = datasets.ImageFolder(root=f"{data_path_cutmix}", transform = imagenet_valid_transform())
val_loader_cutmix = torch.utils.data.DataLoader(val_dataset_cutmix, batch_size=50, shuffle=False, num_workers=4)



#loader_box = {"default" : val_loader_default, "fsmr": val_loader_fsmr, "stylegan" : val_loader_stylegan, "cat": val_loader_cat, "dog" : val_loader_dog, "cutmix" : val_loader_cutmix}
loader_box = {"cutmix" : val_loader_cutmix}

# loss_fn = nn.CrossEntropyLoss().cuda()
model=model.to('cuda')

model.eval()

for loader_name, loader in loader_box.items():
    for batch_idx, (img, target) in enumerate(loader):
        img, target = img.to('cuda'), target.to('cuda')
        with torch.no_grad():
            #IN
            output = model.conv1(img)
            output = model.bn1(output)
            output = model.relu(output)
            output = model.maxpool(output)
            output = model.layer1(output)
            output = model.layer2(output)
            output = model.layer3(output)
            output = model.layer4(output)
            output = model.avgpool(output)

        
            
        out_list = output.squeeze().cpu().detach().numpy() if batch_idx ==0 else np.concatenate((out_list, output.squeeze().cpu().detach().numpy()), axis=0)

    #IN
    out_path = f'./result_npy/IN/'
    if not os.path.isdir(out_path):
        os.mkdir(out_path)
        
    np.save(os.path.join(out_path,f'{loader_name}.npy'), out_list)

 

loader 두개 한번에 : zip 활용

for iter, ((cat,cat_label), (dog,dog_label)) in tqdm(enumerate(zip(val_loader_cat, val_loader_dog)), total = len(val_loader_cat)):
    #if cat.shape[0] != 50 : break
    lam = np.random.beta(1.0, 1.0)  # 베타 분포에서 lam 값을 가져옵나다.
    cat, dog = cat.to('cuda'), dog.to('cuda')
    #import pdb; pdb.set_trace()
    bbx1, bby1, bbx2, bby2 = rand_bbox((3, 256, 256), lam)
    dog[:, :,bbx1:bbx2, bby1:bby2] = cat[:, :,bbx1:bbx2, bby1:bby2]
    
    for i in range(50):
        save_image(dog[i], os.path.join(out_dir, f'img{50*iter+i}.png'))