딥러닝 모델에서의 Quantization이란 부동소수점 형식으로 연산되는 모델의 학습과 추론을 정수형으로 변환하는 것과 같은 방법을 통하여 학습과 추론의 시간을 단축하고 모델의 크기를 줄이는 것이다. 이때, Quantization을 하는 과정에서 정확도가 Quantization을 하기 이전의 모델과 비교했을 때, 크게 떨어져서는 안된다. Pytorch가 패키지 내에서 지원하는 Quantization의 종류로는 1) dynamic quantization, 2) static quantization, 3) quantization aware training이 있는데, 이 중 3)quantization aware training만이 gpu학습을 지원한다. 또한, 위 세가지 방법 중 quantization aware training이 가장 정확도를 유지한다고 한다. Quantization Aware Training을 MNIST 데이터로 실습해보았다.
먼저 Quantization을 할 모델을 정의해준다. QuanStub를 forward 시작하는 부분에 적어주고, DeQuantStub를 forward가 끝나는 부분에 적어주는 것이 특징이다. 이때, LogSoftmax의 경우, 후에 추론 시에 사용할 QuantizedCPU형을 지원하지 않으므로, LogSoftmax이전에 dequantization을 해준다. 파이토치에서 layer별로 지원하는 데이터 형태에 제약이 있는 듯 하다.
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import time
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.quant = torch.quantization.QuantStub()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.relu = torch.nn.ReLU()
self.dropout1= nn.Dropout(0.25)
self.dropout2 = nn.Dropout(0.5)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
self.maxpool = nn.MaxPool2d(2)
self.logsoftmax = torch.nn.LogSoftmax(dim=1)
self.dequant = torch.quantization.DeQuantStub()
def forward(self, x):
x = self.quant(x) # QuanStub를 forward를 시작하는 부분에 적어준다.
x = self.conv1(x)
x = self.relu(x)
x = self.conv2(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.dropout1(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = self.relu(x)
x = self.dropout2(x)
x = self.fc2(x)
x = self.dequant(x) # DeQuanStub를 forward가 끝나는 부분에 적어준다. LogSoftmax의 경우, 후에 추론 시에 사용할 데이터 형태인 QuantizedCPU형을 지원하지 않으므로, LogSofmax이전에 dequantization을 해준다.
x = self.logsoftmax(x)
return x다음으로 CUDA를 이용해서 학습을 진행하고 floating point로 학습한 모델을 quantized integer model로 바꿔준다. Quantized integer model로 바꿔준 후에는, test를 진행한다. 모델의 정확도를 측정하였고 소요시간은 학습부터 추론까지 총 걸린 시간, test를 하기 이전까지 경과 시간, inference를 할 때 걸린 시간으로 구분하여 측정하였다.
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
device = 'cuda'
dataset1 = datasets.MNIST('../data', train=True, download=True,
transform = transform)
dataset2 = datasets.MNIST('../data', train=False, download=True,
transform = transform)
train_loader = torch.utils.data.DataLoader(dataset1, batch_size=64)
test_loader = torch.utils.data.DataLoader(dataset2, batch_size=64)
start = time.time()
model_fp32 = Net()
model_fp32.train() # 아래에 진행될 Quantization Aware Training logic이 작동하기 위해서는 모델을 train 모드로 바꿔줘야 한다고 한다.
model_fp32.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
model_fp32_fused = torch.quantization.fuse_modules(model_fp32, [['conv1', 'relu']])
model_fp32_prepared = torch.quantization.prepare_qat(model_fp32_fused)
model_fp32_prepared = model_fp32_prepared.to("cuda")
optimizer = optim.SGD(model_fp32_prepared.parameters(), lr=0.01, momentum=0.5)
##CUDA를 이용해서 학습한다.
for epoch in range(3):
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model_fp32_prepared(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx & 10 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%]\tLoss:{:.6f}'.format(
epoch, batch_idx*len(data), len(train_loader.dataset),
100.*batch_idx / len(train_loader), loss.item()
))
model_fp32_prepared.eval()
model_int8 = torch.quantization.convert(model_fp32_prepared.to('cpu')) #quantized aware training을 floating point로 수행한 model을 quantized integer model로 바꿔준다.
model_int8.eval()
test_loss = 0
correct = 0
start2 = time.time()
with torch.no_grad():
for data, target in test_loader:
data, target = data.to('cpu'), target.to('cpu') #GPU는 integer형 연산을 지원하지 않으므로 추론 속도를 비교하기 위해서 모델과 data를 모두 cpu로 옮겨줬다.
output = model_int8(data)
test_loss += F.nll_loss(output, target, reduction='sum').item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)
))
end = time.time()
print("test 이전까지 경과 시간(secs):",start2-start)
print("inference를 할 때 걸린 시간(secs):",end-start2)
print("total time elapsed(secs):", (end-start))
Test set: Average loss: 0.0607, Accuracy: 9810/10000 (98%)
test 이전까지 경과 시간(secs): 88.42995023727417
inference를 할 때 걸린 시간(secs): 8.547651767730713
total time elapsed(secs): 96.97760200500488다음으로는 Quantization Aware Training을 하지 않은 일반적인 CNN모델을 정의하고 학습과 추론에 걸리는 시간을 측정하였다. 이때, 추론을 할 때는 위에서 Quantization Aware Training을 한 모델과 동일한 기준으로 비교하기 위해 cpu를 사용하였다.
start = time.time()
model_compared = Net2().to('cuda')
optimizer = optim.SGD(model_compared.parameters(), lr=0.01, momentum=0.5)
for epoch in range(3):
model_compared.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model_compared(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx & 10 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%]\tLoss:{:.6f}'.format(
epoch, batch_idx*len(data), len(train_loader.dataset),
100.*batch_idx / len(train_loader), loss.item()
))
# Inference시의 속도를 Quantization Aware Trained된 모델과 동일한 기준으로 비교하기 위해 모델, 데이터를 cpu로 옮겨줘서 cpu로 연산한다.
model_compared.to('cpu')
model_compared.eval()
test_loss = 0
correct = 0
start2 = time.time()
with torch.no_grad():
for data, target in test_loader:
data, target = data.to('cpu'), target.to('cpu')
output = model_compared(data)
test_loss += F.nll_loss(output, target, reduction='sum').item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)
))
end = time.time()
print("test 이전까지 경과 시간(secs):",start2-start)
print("inference를 할 때 걸린 시간(secs):",end-start2)
print("total time elapsed(secs):", (end-start))
Test set: Average loss: 0.0770, Accuracy: 9759/10000 (98%)
test 이전까지 경과 시간(secs): 57.33148193359375
inference를 할 때 걸린 시간(secs): 9.191223621368408
total time elapsed(secs): 66.52270555496216흥미롭게도, Quantization Aware Training을 했을 때의 모델의 정확도가 더 높다.(둘다 정확도는 98%이지만, Quantization Aware Training의 경우, test데이터 10,000개 중 9,810개를 맞췄고, Quantization을 하지 않은 모델의 경우, 10,000개 중 9,759개를 맞췄다.) 다만, Quantization Aware Training을 진행한 모델의 총 소요시간(96.9초)이 더 길게 나왔다. 이는 test를 하기 전 모델을 Quantization을 하는데 걸린 시간 때문으로 판단된다. 왜냐하면 inference를 하는 데 걸린 시간을 Qunatizaition Aware Training을 진행한 모델이 8.5초로 더 빠르기 때문이다.
