LSTM 单词语言模型上的动态量化（beta）

原文：https://pytorch.org/tutorials/advanced/dynamic_quantization_tutorial.html

作者： James Reed

编辑：Seth Weidman

简介

量化涉及将模型的权重和激活从float转换为int，这可以导致模型尺寸更小，推断速度更快，而对准确率的影响很小。

在本教程中，我们将最简单的量化形式-动态量化应用于基于 LSTM 的下一个单词预测模型，紧紧遵循 PyTorch 示例中的单词语言模型 。

  1. # imports
  2. import os
  3. from io import open
  4. import time
  6. import torch
  7. import torch.nn as nn
  8. import torch.nn.functional as F

1.定义模型

在这里，我们根据词语言模型示例中的模型定义 LSTM 模型架构。

  1. class LSTMModel(nn.Module):
  2.     """Container module with an encoder, a recurrent module, and a decoder."""
  4.     def __init__(self, ntoken, ninp, nhid, nlayers, dropout=0.5):
  5.         super(LSTMModel, self).__init__()
  6.         self.drop = nn.Dropout(dropout)
  7.         self.encoder = nn.Embedding(ntoken, ninp)
  8.         self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)
  9.         self.decoder = nn.Linear(nhid, ntoken)
  11.         self.init_weights()
  13.         self.nhid = nhid
  14.         self.nlayers = nlayers
  16.     def init_weights(self):
  17.         initrange = 0.1
  18.         self.encoder.weight.data.uniform_(-initrange, initrange)
  19.         self.decoder.bias.data.zero_()
  20.         self.decoder.weight.data.uniform_(-initrange, initrange)
  22.     def forward(self, input, hidden):
  23.         emb = self.drop(self.encoder(input))
  24.         output, hidden = self.rnn(emb, hidden)
  25.         output = self.drop(output)
  26.         decoded = self.decoder(output)
  27.         return decoded, hidden
  29.     def init_hidden(self, bsz):
  30.         weight = next(self.parameters())
  31.         return (weight.new_zeros(self.nlayers, bsz, self.nhid),
  32.                 weight.new_zeros(self.nlayers, bsz, self.nhid))

2.加载文本数据

接下来，我们再次根据单词模型示例对预处理，将 Wikitext-2 数据集加载到语料库中。

  1. class Dictionary(object):
  2.     def __init__(self):
  3.         self.word2idx = {}
  4.         self.idx2word = []
  6.     def add_word(self, word):
  7.         if word not in self.word2idx:
  8.             self.idx2word.append(word)
  9.             self.word2idx[word] = len(self.idx2word) - 1
  10.         return self.word2idx[word]
  12.     def __len__(self):
  13.         return len(self.idx2word)
  15. class Corpus(object):
  16.     def __init__(self, path):
  17.         self.dictionary = Dictionary()
  18.         self.train = self.tokenize(os.path.join(path, 'train.txt'))
  19.         self.valid = self.tokenize(os.path.join(path, 'valid.txt'))
  20.         self.test = self.tokenize(os.path.join(path, 'test.txt'))
  22.     def tokenize(self, path):
  23.         """Tokenizes a text file."""
  24.         assert os.path.exists(path)
  25.         # Add words to the dictionary
  26.         with open(path, 'r', encoding="utf8") as f:
  27.             for line in f:
  28.                 words = line.split() + ['<eos>']
  29.                 for word in words:
  30.                     self.dictionary.add_word(word)
  32.         # Tokenize file content
  33.         with open(path, 'r', encoding="utf8") as f:
  34.             idss = []
  35.             for line in f:
  36.                 words = line.split() + ['<eos>']
  37.                 ids = []
  38.                 for word in words:
  39.                     ids.append(self.dictionary.word2idx[word])
  40.                 idss.append(torch.tensor(ids).type(torch.int64))
  41.             ids = torch.cat(idss)
  43.         return ids
  45. model_data_filepath = 'data/'
  47. corpus = Corpus(model_data_filepath + 'wikitext-2')

3.加载预先训练的模型

这是一本有关动态量化的教程，这是在训练模型后应用的一种量化技术。 因此，我们将简单地将一些预训练的权重加载到此模型架构中； 这些权重是通过使用单词语言模型示例中的默认设置训练五个周期而获得的。

  1. ntokens = len(corpus.dictionary)
  3. model = LSTMModel(
  4.     ntoken = ntokens,
  5.     ninp = 512,
  6.     nhid = 256,
  7.     nlayers = 5,
  8. )
  10. model.load_state_dict(
  11.     torch.load(
  12.         model_data_filepath + 'word_language_model_quantize.pth',
  13.         map_location=torch.device('cpu')
  14.         )
  15.     )
  17. model.eval()
  18. print(model)

出：

  1. LSTMModel(
  2.   (drop): Dropout(p=0.5, inplace=False)
  3.   (encoder): Embedding(33278, 512)
  4.   (rnn): LSTM(512, 256, num_layers=5, dropout=0.5)
  5.   (decoder): Linear(in_features=256, out_features=33278, bias=True)
  6. )

现在，我们生成一些文本以确保预先训练的模型能够正常工作-与以前类似，我们在此处遵循

  1. input_ = torch.randint(ntokens, (1, 1), dtype=torch.long)
  2. hidden = model.init_hidden(1)
  3. temperature = 1.0
  4. num_words = 1000
  6. with open(model_data_filepath + 'out.txt', 'w') as outf:
  7.     with torch.no_grad():  # no tracking history
  8.         for i in range(num_words):
  9.             output, hidden = model(input_, hidden)
  10.             word_weights = output.squeeze().div(temperature).exp().cpu()
  11.             word_idx = torch.multinomial(word_weights, 1)[0]
  12.             input_.fill_(word_idx)
  14.             word = corpus.dictionary.idx2word[word_idx]
  16.             outf.write(str(word.encode('utf-8')) + ('\n' if i % 20 == 19 else ' '))
  18.             if i % 100 == 0:
  19.                 print('| Generated {}/{} words'.format(i, 1000))
  21. with open(model_data_filepath + 'out.txt', 'r') as outf:
  22.     all_output = outf.read()
  23.     print(all_output)

出：

  1. | Generated 0/1000 words
  2. | Generated 100/1000 words
  3. | Generated 200/1000 words
  4. | Generated 300/1000 words
  5. | Generated 400/1000 words
  6. | Generated 500/1000 words
  7. | Generated 600/1000 words
  8. | Generated 700/1000 words
  9. | Generated 800/1000 words
  10. | Generated 900/1000 words
  11. b'broadcaster' b'good' b',' b'which' b'provided' b'for' b'a' b'vignettes' b'socially' b'and' b'the' b'FIA' b"'s" b'ad' b'.' b'The' b'state' b'into' b'this' b'position'
  12. b'is' b'in' b'account' b'of' b'a' b'wide' b'Domonia' b'<unk>' b',' b'fallen' b'to' b'for' b'the' b'types' b'of' b'<unk>' b'developers' b'being' b'entertaining' b'.'
  13. b'<eos>' b'The' b'Claus' b'II' b'(' b'The' b'Book' b'of' b'Karnataka' b',' b'2' b'/' b'10' b')' b'was' b'released' b'by' b'British' b'@-@' b'Irish'
  14. b'ruler' b'arriving' b'on' b'the' b'winter' b'of' b'its' b'championship' b'orbit' b'.' b'In' b'early' b'spring' b'roles' b'dismay' b'when' b'he' b'replaced' b'by' b'a'
  15. b'religious' b'park' b',' b'when' b'it' b'features' b'flowers' b'they' b'do' b'populist' b'.' b'temperatures' b'attempted' b'to' b'have' b'trouble' b'met' b',' b'<unk>' b','
  16. b'and' b'karaoke' b'leads' b'to' b'some' b'return' b'up' b'as' b'or' b'seated' b'.' b'The' b'remainder' b'of' b'w' b'voltage' b'contains' b'Allah' b'in' b'the'
  17. b'series' b'to' b'infiltrate' b'disappeared' b'.' b'Though' b'it' b'comes' b'into' b'his' b'Shinnok' b"'s" b'history' b',' b'they' b'may' b'sometimes' b'7' b'@-@' b'April'
  18. b',' b'roughly' b'7' b'%' b'of' b'50' b'mph' b'(' b'4' b'@.@' b'8' b'in' b')' b'while' b'males' b'have' b'put' b'except' b'far' b'as'
  19. b'alkaline' b'@-@' b'up' b'.' b'<eos>' b'Electrical' b'medical' b'rings' b'were' b'always' b'published' b'.' b'<eos>' b'Based' b'on' b'2' b'November' b',' b'Idaho' b'can'
  20. b'be' b'estimated' b'cooking' b'and' b'<unk>' b',' b'while' b'no' b',' b'thin' b'drugs' b'was' b'poor' b'to' b'each' b'area' b'.' b'It' b'has' b'not'
  21. b'campaigned' b'those' b'of' b'the' b'most' b'potent' b'population' b'of' b'leaves' b'in' b'all' b'condition' b',' b'because' b'they' b'were' b'forced' b'to' b'die' b'in'
  22. b'bhandara' b'<unk>' b'that' b'culture' b'.' b'Almost' b'a' b'prose' b'plan' b',' b'there' b'have' b'been' b'only' b'clear' b',' b'it' b'occurs' b'.' b'<eos>'
  23. b'The' b'kakapo' b'was' b'interpreted' b'on' b'1998' b'from' b'1955' b'and' b'played' b'in' b'<unk>' b',' b'Western' b'Asia' b'on' b'0' b'August' b'1966' b','
  24. b'with' b'an' b'additional' b'population' b'that' b'Samuel' b'solemnly' b',' b'Chapman' b'sponsored' b'after' b'a' b'few' b'years' b'.' b'In' b'1990' b',' b'prominent' b'areas'
  25. b'believe' b'that' b'as' b'being' b'an' b'rural' b'planet' b',' b'they' b'is' b'neglected' b'as' b'to' b'be' b'changed' b'.' b'Congress' b'This' b'well' b'"'
  26. b'was' b'run' b'by' b'<unk>' b',' b'Waldemar' b'Greenwood' b'.' b'170' b'have' b'just' b'in' b'place' b',' b'he' b'overruled' b'.' b'The' b'1966' b'race'
  27. b'is' b'a' b'embodies' b'state' b'of' b'Viking' b'or' b'most' b'generation' b',' b'not' b'in' b'the' b'codes' b'of' b'all' b'other' b'alignment' b'musical' b'politicians'
  28. b'.' b'No' b'system' b'have' b'participated' b'on' b'3' b'to' b'9' b'%' b'of' b'any' b'urine' b',' b'with' b'both' b'drawings' b'and' b'significantly' b'towards'
  29. b'his' b'deteriorating' b'and' b'poverty' b'.' b'As' b'a' b'rust' b',' b'contains' b'other' b'compositions' b'that' b'must' b'be' b'beneficial' b'by' b'overnight' b'or' b'fluid'
  30. b',' b'u' b'organizations' b'can' b'seek' b'mild' b'late' b'down' b'on' b'a' b'broadside' b'and' b'leads' b'to' b'its' b'cycle' b'.' b'For' b'example' b','
  31. b'1137' b',' b'snowmelt' b'and' b'<unk>' b'\xe2\x80\x94' b'a' b'variety' b'of' b'dealt' b';' b'Species' b'(' b'with' b'a' b'reduction' b'of' b'prohibitions' b')' b','
  32. b'<unk>' b'exploration' b',' b'<unk>' b'an' b'fuel' b'eye' b'of' b'purple' b'trees' b',' b'was' b'shown' b'west' b'.' b'chased' b'Jack' b'of' b'claws' b','
  33. b'his' b'vertex' b'states' b'that' b'they' b',' b'in' b'1922' b',' b'was' b'killed' b'.' b'<eos>' b'There' b'have' b'been' b'official' b'concerns' b'of' b'Boat'
  34. b'Kerry' b'including' b'L\xc3\xaa' b'\xe3\x80\x89' b'and' b'<unk>' b'A' b'Forest' b',' b'"' b'<unk>' b',' b'because' b'<unk>' b',' b'and' b'sometimes' b'encounters' b'like' b'I'
  35. b"'ve" b'been' b'<unk>' b'.' b'"' b'<unk>' b'Hunter' b'pathway' b'writes' b'it' b'entering' b'the' b'second' b'.' b'The' b'kakapo' b'is' b'gems' b'used' b'after'
  36. b'died' b'from' b'two' b'games' b'in' b'six' b'<unk>' b',' b'her' b'feature' b'and' b'called' b'"' b'mercenaries' b'"' b',' b'which' b'supported' b'by' b'the'
  37. b'Selective' b'Race' b'.' b'"' b'<eos>' b'Bono' b'Dutch' b'struggles' b'to' b'the' b'species' b'<unk>' b',' b'especially' b'crusaders' b'I' b'lives' b'process' b',' b'but'
  38. b'Constantin' b'approximate' b'and' b'character' b'or' b'so' b'.' b'There' b'have' b'numerous' b'pale' b'dioceses' b'as' b'a' b'resistant' b';' b'the' b'Inn' b'Comic' b'@-@'
  39. b'white' b'individuals' b',' b'its' b'flat' b',' b'<unk>' b'and' b'correct' b',' b'in' b'which' b'they' b'felt' b'.' b'In' b'the' b'arms' b',' b'the'
  40. b'original' b'occasion' b'about' b'Spanish' b'sites' b'all' b'(' b'millionaire' b'lay' b';' b'or' b'160' b'@-@' b'mosquitoes' b')' b'v' b'<unk>' b'(' b'c' b')'
  41. b'.' b'The' b'bird' b'is' b'extremely' b'paved' b',' b'and' b'they' b'are' b'claimed' b'to' b'wedding' b'the' b'<unk>' b'of' b'Excellence' b',' b'and' b'an'
  42. b'extinct' b'composite' b',' b'cute' b'outside' b'<unk>' b'.' b'This' b'may' b'be' b'seen' b'by' b'the' b'Seer' b'that' b'Tempest' b'"' b'comes' b'"' b'over'
  43. b'a' b'bright' b'judicial' b'guitar' b',' b'which' b'describes' b',' b'and' b'tend' b'to' b'be' b'seen' b'.' b'<eos>' b'<eos>' b'=' b'=' b'Conservation' b'for'
  44. b'contraception' b'=' b'=' b'<eos>' b'<eos>' b'Grieco' b'Island' b'is' b'a' b'eventually' b'scale' b'word' b'to' b'a' b'tropical' b'storm' b',' b'based' b'in' b'a'
  45. b'pre' b'\xe2\x80\x93' b'9' b'lead' b',' b'a' b'forces' b'after' b'a' b'additional' b',' b'grey' b'substance' b',' b'Metro' b',' b'background' b',' b'and' b'cooperate'
  46. b'with' b'its' b'overly' b'overview' b',' b'so' b'the' b'heaviest' b'route' b',' b'and' b'\xc2\xb3' b'.' b'portion' b'may' b'occur' b'this' b'other' b'up' b'an'
  47. b'<unk>' b'break' b',' b'then' b'or' b'deep' b'distinct' b'or' b'female' b'offspring' b',' b'but' b'even' b'understand' b'.' b'Following' b'God' b'(' b'no' b'nervous'
  48. b'image' b'from' b'complaints' b')' b',' b'the' b'player' b'represents' b'three' b'or' b'over' b'9' b'\xc2\xb0' b'large' b'(' b'five' b'weeks' b'of' b'many' b'cats'
  49. b')' b',' b'as' b'it' b'targets' b'for' b'the' b'second' b'female' b'together' b'.' b'159' b',' b'it' b'also' b'spend' b'bold' b'markets' b'and' b'its'
  50. b'players' b'powers' b',' b'dubbed' b'those' b'of' b'lengths' b'.' b'Most' b'are' b'arrow' b'could' b'be' b'noticed' b'involving' b'they' b'fall' b'.' b'On' b'FAU'
  51. b"'s" b'only' b'lifetime' b'she' b'treated' b'or' b'their' b'apparent' b'soaring' b'proposition' b'has' b'5th' b'of' b'those' b'eye' b',' b'but' b'knows' b'in' b'a'
  52. b'<unk>' b'Network' b';' b'which' b'of' b'that' b'reality' b'or' b'artificial' b'when' b'struggling' b'Bungie' b'is' b'successful' b'.' b'The' b'<unk>' b'sound' b'of' b'frontier'
  53. b'ahead' b'for' b'damage' b'came' b'on' b',' b'so' b'the' b'first' b'series' b'funded' b'by' b'its' b'bowls' b',' b'a' b'chant' b'.' b'They' b'may'
  54. b'be' b'used' b'Pongsak' b'or' b'occasionally' b'protected' b'them' b'.' b'Fingal' b'cylindrical' b'conspired' b'on' b'a' b'variety' b'of' b'prey' b',' b'<unk>' b',' b'Zach'
  55. b',' b'and' b'young' b'possessing' b'Westland' b'valleys' b'.' b'Otherwise' b',' b'I' b'do' b'at' b'them' b'in' b'first' b'@-@' b'season' b'woodland' b',' b'where'
  56. b'they' b'weighed' b'them' b'to' b'correct' b'a' b'list' b'of' b'other' b'birds' b'.' b'Another' b'theme' b'where' b'or' b',' b'it' b'is' b'a' b'appropriate'
  57. b'source' b',' b'this' b'competed' b'in' b'integral' b'Waiouru' b'alone' b',' b'the' b'pathways' b'under' b'Aravind' b',' b'and' b'others' b',' b'instead' b'of' b'westward'
  58. b',' b'as' b'they' b'are' b'quarters' b'and' b'caused' b'in' b'males' b'.' b'Once' b'selective' b'centered' b',' b'they' b'threats' b'were' b'Zuniceratops' b'.' b'Although'
  59. b'the' b'most' b'spots' b'replication' b'became' b'a' b'fragile' b'pointer' b'(' b'a' b'pair' b'of' b'<unk>' b')' b',' b'strongly' b'"' b'mammals' b'"' b','
  60. b'which' b'give' b'Powderfinger' b'to' b'persecution' b'.' b'Other' b'conifers' b'but' b'even' b'only' b'swallow' b'so' b'every' b'symbols' b'of' b'Manders' b',' b'in' b'massive'

它不是 GPT-2，但看起来该模型已开始学习语言结构！

我们几乎准备好演示动态量化。 我们只需要定义一些辅助函数：

  1. bptt = 25
  2. criterion = nn.CrossEntropyLoss()
  3. eval_batch_size = 1
  5. # create test data set
  6. def batchify(data, bsz):
  7.     # Work out how cleanly we can divide the dataset into bsz parts.
  8.     nbatch = data.size(0) // bsz
  9.     # Trim off any extra elements that wouldn't cleanly fit (remainders).
  10.     data = data.narrow(0, 0, nbatch * bsz)
  11.     # Evenly divide the data across the bsz batches.
  12.     return data.view(bsz, -1).t().contiguous()
  14. test_data = batchify(corpus.test, eval_batch_size)
  16. # Evaluation functions
  17. def get_batch(source, i):
  18.     seq_len = min(bptt, len(source) - 1 - i)
  19.     data = source[i:i+seq_len]
  20.     target = source[i+1:i+1+seq_len].reshape(-1)
  21.     return data, target
  23. def repackage_hidden(h):
  24.   """Wraps hidden states in new Tensors, to detach them from their history."""
  26.   if isinstance(h, torch.Tensor):
  27.       return h.detach()
  28.   else:
  29.       return tuple(repackage_hidden(v) for v in h)
  31. def evaluate(model_, data_source):
  32.     # Turn on evaluation mode which disables dropout.
  33.     model_.eval()
  34.     total_loss = 0.
  35.     hidden = model_.init_hidden(eval_batch_size)
  36.     with torch.no_grad():
  37.         for i in range(0, data_source.size(0) - 1, bptt):
  38.             data, targets = get_batch(data_source, i)
  39.             output, hidden = model_(data, hidden)
  40.             hidden = repackage_hidden(hidden)
  41.             output_flat = output.view(-1, ntokens)
  42.             total_loss += len(data) * criterion(output_flat, targets).item()
  43.     return total_loss / (len(data_source) - 1)

4.测试动态量化

最后，我们可以在模型上调用torch.quantization.quantize_dynamic！ 特别，

  • 我们指定我们希望对模型中的nn.LSTM和nn.Linear模块进行量化
  • 我们指定希望将权重转换为int8值
  1. import torch.quantization
  3. quantized_model = torch.quantization.quantize_dynamic(
  4.     model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
  5. )
  6. print(quantized_model)

出：

  1. LSTMModel(
  2.   (drop): Dropout(p=0.5, inplace=False)
  3.   (encoder): Embedding(33278, 512)
  4.   (rnn): DynamicQuantizedLSTM(512, 256, num_layers=5, dropout=0.5)
  5.   (decoder): DynamicQuantizedLinear(in_features=256, out_features=33278, dtype=torch.qint8, qscheme=torch.per_tensor_affine)
  6. )

该模型看起来相同； 这对我们有什么好处？ 首先，我们看到模型尺寸显着减小：

  1. def print_size_of_model(model):
  2.     torch.save(model.state_dict(), "temp.p")
  3.     print('Size (MB):', os.path.getsize("temp.p")/1e6)
  4.     os.remove('temp.p')
  6. print_size_of_model(model)
  7. print_size_of_model(quantized_model)

出：

  1. Size (MB): 113.945726
  2. Size (MB): 79.739984

其次，我们看到了更快的推理时间，而评估损失没有差异：

注意：由于量化模型运行单线程，因此用于单线程比较的线程数为 1。

  1. torch.set_num_threads(1)
  3. def time_model_evaluation(model, test_data):
  4.     s = time.time()
  5.     loss = evaluate(model, test_data)
  6.     elapsed = time.time() - s
  7.     print('''loss: {0:.3f}\nelapsed time (seconds): {1:.1f}'''.format(loss, elapsed))
  9. time_model_evaluation(model, test_data)
  10. time_model_evaluation(quantized_model, test_data)

出：

  1. loss: 5.167
  2. elapsed time (seconds): 251.3
  3. loss: 5.168
  4. elapsed time (seconds): 166.3

在没有量化的情况下在 MacBook Pro 上本地运行此操作，推理大约需要 200 秒，而量化则只需大约 100 秒。

总结

动态量化可能是减小模型大小的简单方法，而对精度的影响有限。

谢谢阅读！ 与往常一样，我们欢迎您提供反馈，因此，如果有任何问题，请在这里创建一个 ISSUE。

脚本的总运行时间：（7 分钟 3.126 秒）

下载 Python 源码：dynamic_quantization_tutorial.py

下载 Jupyter 笔记本：dynamic_quantization_tutorial.ipynb

由 Sphinx 画廊生成的画廊