딥러닝 모델을 학습시키다 보면 항상 vram의 압박에 시달리게 된다. 특히 최근 막대한 크기의 모델들이 등장해 이런 압박은 더 심해지기도 한다.
한편, 일반 사용자용 그래픽 카드 중 최상위인 Nvidia 2080ti조차도 vram이 겨우 11GB밖에 되지 않아 거대한 모델을 Fine-tuning 하는 것조차 굉장히 작은 배치사이즈로 학습시켜야 한다.
Google Colab에서 제공하는 TPU는 tpu v3-8 모델로 총 128GB의 메모리를 가지고 있어, 상대적으로 큰 모델과 배치사이즈를 이용해 학습할 수 있다. (tpu v3 하나는 16GB의 HBM 메모리를 가지고 있고, tpu v3-8은 8개의 코어로 총 128GB의 메모리를 가진다.)
PyTorch에서는 Pytorch/XLA 프로젝트를 통해 PyTorch에서도 TPU를 통한 학습을 할 수 있도록 컴파일러를 제공하고 있고, colab에 해당 패키지를 설치하면 TPU를 곧바로 사용할 수 있다.
# See https://pytorch.org/docs/stable/torchvision/models.html for normalization # Pre-trained TorchVision models expect RGB (3 x H x W) images # H and W should be >= 224 # Loaded into [0, 1] and normalized as follows: normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) to_rgb = transforms.Lambda(lambda image: image.convert('RGB')) resize = transforms.Resize((224, 224)) my_transform = transforms.Compose([resize, to_rgb, transforms.ToTensor(), normalize])
torchvision.transforms를 통해 데이터셋을 이미지넷 기반으로 Normalize하고 & RGB & 244x244 사이즈로 리사이징 하는 처리를 해줄 수 있다.
# Note: this will take 4-5 minutes to run. num_epochs = 1 loss_fn = torch.nn.CrossEntropyLoss() optimizer = torch.optim.Adam(net.parameters())
# Ensures network is in train mode net.train()
start_time = time.time() for epoch in range(num_epochs): for data, targets in iter(train_loader): # Sends data and targets to device data = data.to(device) targets = targets.to(device)
# Acquires the network's best guesses at each class results = net(data)
# Computes loss loss = loss_fn(results, targets)
# Updates model optimizer.zero_grad() loss.backward() xm.optimizer_step(optimizer, barrier=True) # 이부분이 TPU 쓸때 필요한 코드!!
elapsed_time = time.time() - start_time print ("Spent ", elapsed_time, " seconds training for ", num_epochs, " epoch(s) on a single core.")
위 코드 중 25번째 줄의 xm.optimizer_step(optimizer, barrier=True) 부분이 TPU를 사용하기 위한 코드이다. (GPU 코드에 겨우 한줄 추가!)
TPU 코어 FULL 활용하기
앞서 다룬 과정에서는 TPU 8core 중 1core만을 사용한다. 한편, TPU 1개에 들어있는 8개의 코어 전체를 사용하면 보다 빠른 학습이 가능하다.
import torch import torch_xla import torch_xla.core.xla_model as xm import torch_xla.distributed.xla_multiprocessing as xmp
# "Map function": acquires a corresponding Cloud TPU core, creates a tensor on it, # and prints its core defsimple_map_fn(index, flags): # Sets a common random seed - both for initialization and ensuring graph is the same torch.manual_seed(1234)
# Acquires the (unique) Cloud TPU core corresponding to this process's index device = xm.xla_device()
# Creates a tensor on this process's device t = torch.randn((2, 2), device=device)
print("Process", index ,"is using", xm.xla_real_devices([str(device)])[0])
# Spawns eight of the map functions, one for each of the eight cores on # the Cloud TPU flags = {} # Note: Colab only supports start_method='fork' xmp.spawn(simple_map_fn, args=(flags,), nprocs=8, start_method='fork')
TPU 여러 코어를 사용하기 위해서는 torch_xla.distributed.xla_multiprocessing 을 통해 프로세스를 N개 띄우는 방식으로 진행한다.
단, 앞서 진행한 코드는 model과 data 로드 부분이 모두 각자의 코드로 나와 Colab 인스턴스의 CPU/Ram에서 진행되어 코드 내의 변수를 .to(device) 를 사용해 GPU나 TPU로 보낼 수 있지만, TPU의 여러 코어를 사용할때는 하나의 함수 내에 model과 data 모두를 넣고 진행해야 한다.
import torchvision from torchvision import datasets import torchvision.transforms as transforms import torch_xla.distributed.parallel_loader as pl import time
train_start = time.time() for epoch in range(flags['num_epochs']): para_train_loader = pl.ParallelLoader(train_loader, [device]).per_device_loader(device) for batch_num, batch in enumerate(para_train_loader): data, targets = batch
output = net(data)
loss = loss_fn(output, targets)
optimizer.zero_grad() loss.backward()
xm.optimizer_step(optimizer) # ParallelLoader 쓸때는 barrier=True 필요 없음
