※ PyTorch에서 특정한 학습 데이터 세트를 불러와 변형한 뒤에 저장하기 ※
흔히 PyTorch를 이용해 학습을 진행하다 보면, 학습 데이터 세트에 변형을 가하는 경우가 많다. 이미지를 회전시키거나(rotation), 적대적 노이즈를 섞거나(adversarial attack), 이미지에 패치를 붙이거나(patching) 등의 작업을 수행하는 경우가 많다. 그렇다면 하나의 데이터 세트를 불러와 다양한 전처리를 수행한 뒤에, 전처리된 데이터를 저장해 놓고 나중에 사용하고자 한다면 어떻게 하면 될까?
아래는 학습 데이터를 처리한 뒤에 저장하는 예시이다. 우리는 흔히 학습 데이터 세트를 데이터 로더(data loader)를 이용해 배치(batch) 단위로 불러와 학습한다. 이때 필자는 Epoch별로 서로 다른 image transformation을 적용하고, 각 Epoch마다 변형된 데이터 세트를 저장하고자 했다. 그래서 다음과 같이 코드를 작성했고, Epoch마다 전체 학습 데이터 세트를 하나의 Tensor 형태로 통째로 저장(save)하기 때문에, 많은 메모리 용량이 소요된다. (CIFAR-10 기준으로 Epoch당 약 600MB) 그래도 하나의 Tensor를 통째로 저장한다는 측면에서 저장하는 것 자체는 그다지 많은 시간이 걸리지 않았다.
epoch = 0 # 현재의 Epoch 명시
save = True # 저장할 것인지
saved_images = None
saved_labels = None
for batch_idx, (inputs, targets) in enumerate(data_loader):
inputs, targets = inputs.cuda(), targets.cuda()
"""
자신이 원하는 대로 데이터 세트를 처리하기
- 처리된 입력 데이터: processed_inputs
- 처리된 정답 레이블: processed_targets
"""
# 각 배치(batch)의 데이터를 연결해 하나의 Tensor에 이어 붙이기
if save:
if saved_images == None:
saved_images = processed_inputs
saved_labels = processed_targets
else:
saved_images = torch.cat((saved_images, processed_inputs), 0)
saved_labels = torch.cat((saved_labels, processed_targets), 0)
# 하나의 Tensor 형태로 처리된 데이터 세트 저장하기(epoch별로 저장)
if save:
path = root_path + 'train_dataset/' + str(epoch) + '/'
if not os.path.isdir(path):
os.makedirs(path)
torch.save(saved_images.cpu(), path + 'processed_inputs')
torch.save(saved_labels.cpu(), path + 'processed_targets')
나중에 데이터를 불러올 때는 다음과 같이 torch.utils.data.Dataset 라이브러리를 사용한다. 필자는 단순히 CIFAR-10 학습 데이터 세트에 대해서 만들었기에, 데이터 세트의 크기(length)도 50,000으로 넣었다.
class CustomDataset(torch.utils.data.Dataset):
def __init__(self, images_path, labels_path):
self.images = torch.load(images_path)
self.labels = torch.load(labels_path)
self.length = 50000
def __getitem__(self, index):
image_tensor = self.images[index]
label_tensor = self.labels[index]
return image_tensor, label_tensor
def __len__(self):
return self.length
결과적으로 저장된 데이터 세트를 읽어오는 코드 예시는 다음과 같다.
images_path = root_path + 'train_dataset/' + str(prev_epoch) + '/processed_inputs'
labels_path = root_path + 'train_dataset/' + str(prev_epoch) + '/processed_targets'
custom_dataset = CustomDataset(images_path, labels_path)
custom_loader = torch.utils.data.DataLoader(custom_dataset, batch_size=batch_size, shuffle=False, num_workers=16)
※ 발생할 수 있는 문제 1 ※
참고로 필자의 경우에는, 처음에 저장했던 Tensor를 다시 불러올 때 다음과 같은 CUDA 오류(error)가 발생했다.
"RuntimeError: CUDA error: device-side assert triggered"
당시에 문제의 원인을 생각해 보니, 필자는 데이터를 불러와 학습을 진행할 때 DataParallel을 사용하고 있고, Tensor 데이터를 저장할 때는 Tensor를 CPU로 보내지 않고, 단순히 GPU에 올라가 있는 상태에서 Tensor 데이터를 save()했다. 그래서 저장하기 전에 cpu()를 호출하여 CPU로 보낸 상태에서 Tensor를 저장했더니 오류가 해결되었다.
※ 발생할 수 있는 문제 2 ※
필자가 처음에 Tensor를 저장할 때는 단순히 이미지를 한 장씩 쪼개서 저장했다. 하지만 필자가 생각했던 것보다 저장되는 Tensor의 크기가 더욱 컸다. 3 X 32 X 32의 크기이기 때문에, 한 장의 이미지라는 점에서 하나당 약 12KB(3,072 X 4 Bytes)를 예상했기 때문이다. 그래서 결과적으로 이미지를 한 장씩 저장하지 않고, 60,000개의 데이터를 전부 하나의 Tensor에 concat으로 (concatenation)을 해서 저장했다. 이렇게 했을 때는 (3,072 X 50,000 X 4 Bytes) = 약 600MB의 크기로 데이터 세트가 저장되었다.