Merge pull request #69 from marian-nmt/quant-test

Regression tests for model quantization

13 files changed, 336 insertions, 0 deletions

diff --git a/tests/training/features/quantized-model/.gitignore b/tests/training/features/quantized-model/.gitignore
new file mode 100644
index 0000000..08afa18
--- /dev/null
+++ b/tests/training/features/quantized-model/.gitignore
@@ -0,0 +1 @@
+train
diff --git a/tests/training/features/quantized-model/model_centers.expected b/tests/training/features/quantized-model/model_centers.expected
new file mode 100644
index 0000000..57380ae
--- /dev/null
+++ b/tests/training/features/quantized-model/model_centers.expected
@@ -0,0 +1,51 @@
+Tensor decoder_W_comb_att unique centers: [-0.17677179 -0.11784786 -0.05892393 -0.          0.05892393  0.11784786
+  0.17677179]
+Tensor decoder_Wc_att unique centers: [-0.15336949 -0.10224632 -0.05112316 -0.          0.05112316  0.10224632
+  0.15336949]
+Tensor Wemb_dec unique centers: [-0.32046145 -0.21364096 -0.10682048  0.          0.10682048  0.21364096
+  0.32046145]
+Tensor decoder_U unique centers: [-0.17687811 -0.11791874 -0.05895937 -0.          0.05895937  0.11791874
+  0.17687811]
+Tensor decoder_Ux unique centers: [-0.21770547 -0.14513698 -0.07256849  0.          0.07256849  0.14513698
+  0.21770547]
+Tensor decoder_W unique centers: [-0.19397542 -0.12931694 -0.06465847 -0.          0.06465847  0.12931694
+  0.19397542]
+Tensor decoder_Wx unique centers: [-0.25329626 -0.16886416 -0.08443208 -0.          0.08443208  0.16886416
+  0.25329626]
+Tensor decoder_U_nl unique centers: [-0.17696194 -0.11797463 -0.05898732  0.          0.05898732  0.11797463
+  0.17696194]
+Tensor decoder_Ux_nl unique centers: [-0.21896881 -0.14597921 -0.07298961  0.          0.07298961  0.14597921
+  0.21896881]
+Tensor decoder_Wc unique centers: [-0.15324192 -0.10216128 -0.05108064  0.          0.05108064  0.10216128
+  0.15324192]
+Tensor decoder_Wcx unique centers: [-0.18192002 -0.12128001 -0.06064001 -0.          0.06064001  0.12128001
+  0.18192002]
+Tensor ff_logit_prev_W unique centers: [-0.32183957 -0.2145597  -0.10727985 -0.          0.10727985  0.2145597
+  0.32183957]
+Tensor ff_logit_lstm_W unique centers: [-0.25455362 -0.16970241 -0.08485121  0.          0.08485121  0.16970241
+  0.25455362]
+Tensor ff_logit_ctx_W unique centers: [-0.19867198 -0.13244799 -0.06622399 -0.          0.06622399  0.13244799
+  0.19867198]
+Tensor decoder_ff_logit_l2_Wt unique centers: [-0.36124557 -0.24083039 -0.1204152   0.          0.1204152   0.24083039
+  0.36124557]
+Tensor ff_state_W unique centers: [-0.17704961 -0.11803307 -0.05901653  0.          0.05901653  0.11803307
+  0.17704961]
+Tensor Wemb unique centers: [-0.31208774 -0.20805849 -0.10402925  0.          0.10402925  0.20805849
+  0.31208774]
+Tensor encoder_U unique centers: [-0.17686225 -0.11790817 -0.05895409  0.          0.05895409  0.11790817
+  0.17686225]
+Tensor encoder_Ux unique centers: [-0.21824732 -0.14549822 -0.07274911  0.          0.07274911  0.14549822
+  0.21824732]
+Tensor encoder_W unique centers: [-0.19403435 -0.12935624 -0.06467812  0.          0.06467812  0.12935624
+  0.19403435]
+Tensor encoder_Wx unique centers: [-0.25213736 -0.16809157 -0.08404578 -0.          0.08404578  0.16809157
+  0.25213736]
+Tensor encoder_r_U unique centers: [-0.17699143 -0.11799429 -0.05899715  0.          0.05899715  0.11799429
+  0.17699143]
+Tensor encoder_r_Ux unique centers: [-0.21971346 -0.14647564 -0.07323782 -0.          0.07323782  0.14647564
+  0.21971346]
+Tensor encoder_r_W unique centers: [-0.19410282 -0.12940188 -0.06470094  0.          0.06470094  0.12940188
+  0.19410282]
+Tensor encoder_r_Wx unique centers: [-0.25225359 -0.16816907 -0.08408453 -0.          0.08408453  0.16816907
+  0.25225359]
+Tensor decoder_c_tt unique centers: []
diff --git a/tests/training/features/quantized-model/quantized-log4bit.expected b/tests/training/features/quantized-model/quantized-log4bit.expected
new file mode 100644
index 0000000..1de1b38
--- /dev/null
+++ b/tests/training/features/quantized-model/quantized-log4bit.expected
@@ -0,0 +1,10 @@
+225.47087097
+245.67460632
+232.80288696
+229.51301575
+217.59494019
+207.93991089
+201.96940613
+194.00720215
+197.69416809
+199.42311096
diff --git a/tests/training/features/quantized-model/quantized-opt.expected b/tests/training/features/quantized-model/quantized-opt.expected
new file mode 100644
index 0000000..12f21d3
--- /dev/null
+++ b/tests/training/features/quantized-model/quantized-opt.expected
@@ -0,0 +1,10 @@
+225.11299133
+243.58525085
+229.45321655
+224.28414917
+212.65376282
+204.06687927
+197.81901550
+190.08296204
+193.72265625
+195.21139526
diff --git a/tests/training/features/quantized-model/quantized-with-bias.expected b/tests/training/features/quantized-model/quantized-with-bias.expected
new file mode 100644
index 0000000..4f56c41
--- /dev/null
+++ b/tests/training/features/quantized-model/quantized-with-bias.expected
@@ -0,0 +1,10 @@
+225.10006714
+243.58285522
+229.47399902
+224.31018066
+212.68742371
+204.09915161
+197.85253906
+190.12380981
+193.74653625
+195.23658752
diff --git a/tests/training/features/quantized-model/quantized.expected b/tests/training/features/quantized-model/quantized.expected
new file mode 100644
index 0000000..17620ec
--- /dev/null
+++ b/tests/training/features/quantized-model/quantized.expected
@@ -0,0 +1,10 @@
+225.10929871
+243.58345032
+229.45071411
+224.28813171
+212.65242004
+204.06596375
+197.81690979
+190.08915710
+193.72299194
+195.20808411
diff --git a/tests/training/features/quantized-model/setup.sh b/tests/training/features/quantized-model/setup.sh
new file mode 100644
index 0000000..71af484
--- /dev/null
+++ b/tests/training/features/quantized-model/setup.sh
@@ -0,0 +1,3 @@
+test -f $MRT_DATA/europarl.de-en/corpus.bpe.en || exit 1
+test -f $MRT_DATA/europarl.de-en/corpus.bpe.de || exit 1
+
diff --git a/tests/training/features/quantized-model/test_quant_centers.sh b/tests/training/features/quantized-model/test_quant_centers.sh
new file mode 100644
index 0000000..22dd863
--- /dev/null
+++ b/tests/training/features/quantized-model/test_quant_centers.sh
@@ -0,0 +1,31 @@
+#!/bin/bash -x
+
+#####################################################################
+# SUMMARY: Make sure that the resulting model is in quantized form
+# AUTHOR: afaji
+#####################################################################
+
+# Exit on error
+set -e
+
+PREFIX=test-center
+
+# Remove old artifacts and create working directory
+rm -rf train
+mkdir -p train
+
+# Train an 8-bits model
+$MRT_MARIAN/marian \
+    --no-shuffle --seed 1111 --dim-emb 32 --dim-rnn 64 --mini-batch 32 --maxi-batch 1 --maxi-batch-sort none --learn-rate 0.1 --optimizer sgd \
+    -m train/model.npz -t $MRT_DATA/europarl.de-en/corpus.bpe.{en,de} -v train/vocab.en.yml train/vocab.de.yml \
+    --cost-type cross-entropy --sync-sgd --after-batches 10 --disp-freq 2 --quantize-bits 3
+
+# Check if files exist
+test -e train/model.npz
+
+# make sure that the resulting model has no more than 256 different values (i.e. quantized)
+$MRT_TOOLS/check-model-unique-vals.py train/model.npz -b 3 --print_centers -o model_centers.out
+$MRT_TOOLS/diff-nums.py model_centers.out model_centers.expected -o model_centers.diff --numpy
+
+# Exit with success code
+exit 0
diff --git a/tests/training/features/quantized-model/test_quantmodel.sh b/tests/training/features/quantized-model/test_quantmodel.sh
new file mode 100644
index 0000000..8b55697
--- /dev/null
+++ b/tests/training/features/quantized-model/test_quantmodel.sh
@@ -0,0 +1,36 @@
+#!/bin/bash -x
+
+#####################################################################
+# SUMMARY: Train a quantized marian model
+# AUTHOR: afaji
+#####################################################################
+
+# Exit on error
+set -e
+
+PREFIX=quantized
+
+# Remove old artifacts and create working directory
+rm -rf train $PREFIX.{log,out,diff}
+mkdir -p train
+
+# Train an 8-bits model
+$MRT_MARIAN/marian \
+    --no-shuffle --seed 1111 --dim-emb 32 --dim-rnn 64 --mini-batch 32 --maxi-batch 1 --maxi-batch-sort none --learn-rate 0.1 --optimizer sgd \
+    -m train/model.npz -t $MRT_DATA/europarl.de-en/corpus.bpe.{en,de} -v train/vocab.en.yml train/vocab.de.yml \
+    --cost-type cross-entropy --sync-sgd --after-batches 100 --disp-freq 10 --quantize-bits 8 \
+    --log $PREFIX.log
+
+# Check if files exist
+test -e train/model.npz
+test -e $PREFIX.log
+
+# Compare the current output with the expected output
+cat $PREFIX.log | $MRT_TOOLS/extract-costs.sh > $PREFIX.out
+$MRT_TOOLS/diff-nums.py $PREFIX.out $PREFIX.expected -o $PREFIX.diff
+
+# make sure that the resulting model has no more than 256 different values (i.e. quantized)
+$MRT_TOOLS/check-model-unique-vals.py train/model.npz -b 8
+
+# Exit with success code
+exit 0
diff --git a/tests/training/features/quantized-model/test_quantmodel_log.sh b/tests/training/features/quantized-model/test_quantmodel_log.sh
new file mode 100644
index 0000000..f79809b
--- /dev/null
+++ b/tests/training/features/quantized-model/test_quantmodel_log.sh
@@ -0,0 +1,36 @@
+#!/bin/bash -x
+
+#####################################################################
+# SUMMARY: Train a quantized marian model
+# AUTHOR: afaji
+#####################################################################
+
+# Exit on error
+set -e
+
+PREFIX=quantized-log4bit
+
+# Remove old artifacts and create working directory
+rm -rf train $PREFIX.{log,out,diff}
+mkdir -p train
+
+# Train an 8-bits model
+$MRT_MARIAN/marian \
+    --no-shuffle --seed 1111 --dim-emb 32 --dim-rnn 64 --mini-batch 32 --maxi-batch 1 --maxi-batch-sort none --learn-rate 0.1 --optimizer sgd \
+    -m train/model.npz -t $MRT_DATA/europarl.de-en/corpus.bpe.{en,de} -v train/vocab.en.yml train/vocab.de.yml \
+    --cost-type cross-entropy --sync-sgd --after-batches 100 --disp-freq 10 --quantize-bits 4 --quantize-log-based --quantize-optimization-steps 3 \
+    --log $PREFIX.log
+
+# Check if files exist
+test -e train/model.npz
+test -e $PREFIX.log
+
+# Compare the current output with the expected output
+cat $PREFIX.log | $MRT_TOOLS/extract-costs.sh > $PREFIX.out
+$MRT_TOOLS/diff-nums.py $PREFIX.out $PREFIX.expected -o $PREFIX.diff
+
+# make sure that the resulting model has no more than 256 different values (i.e. quantized)
+$MRT_TOOLS/check-model-unique-vals.py train/model.npz -b 4
+
+# Exit with success code
+exit 0
diff --git a/tests/training/features/quantized-model/test_quantmodel_with_bias.sh b/tests/training/features/quantized-model/test_quantmodel_with_bias.sh
new file mode 100644
index 0000000..de14ffb
--- /dev/null
+++ b/tests/training/features/quantized-model/test_quantmodel_with_bias.sh
@@ -0,0 +1,43 @@
+#!/bin/bash -x
+
+#####################################################################
+# SUMMARY: Train a quantized marian model
+# AUTHOR: afaji
+#####################################################################
+
+# Exit on error
+set -e
+
+PREFIX=quantized-with-bias
+
+# Remove old artifacts and create working directory
+rm -rf train $PREFIX.{log,out,diff}
+mkdir -p train
+
+# training with quantized bias is tricky, so we start by training a normal model first before finetuning it to the quantized space.
+$MRT_MARIAN/marian \
+    --no-shuffle --seed 1111 --dim-emb 32 --dim-rnn 64 --mini-batch 32 --maxi-batch 1 --maxi-batch-sort none --learn-rate 0.1 --optimizer sgd \
+    -m train/model.npz -t $MRT_DATA/europarl.de-en/corpus.bpe.{en,de} -v train/vocab.en.yml train/vocab.de.yml \
+    --cost-type cross-entropy --sync-sgd --after-batches 20 --disp-freq 10 \
+    --log $PREFIX.log
+
+# Train an 8-bits model
+$MRT_MARIAN/marian \
+    --no-shuffle --seed 1111 --dim-emb 32 --dim-rnn 64 --mini-batch 32 --maxi-batch 1 --maxi-batch-sort none --learn-rate 0.1 --optimizer sgd \
+    -m train/model.npz -t $MRT_DATA/europarl.de-en/corpus.bpe.{en,de} -v train/vocab.en.yml train/vocab.de.yml \
+    --cost-type cross-entropy --sync-sgd --after-batches 100 --disp-freq 10 --quantize-bits 8 --quantize-biases \
+    --log $PREFIX.log
+
+# Check if files exist
+test -e train/model.npz
+test -e $PREFIX.log
+
+# Compare the current output with the expected output
+cat $PREFIX.log | $MRT_TOOLS/extract-costs.sh > $PREFIX.out
+$MRT_TOOLS/diff-nums.py $PREFIX.out $PREFIX.expected -o $PREFIX.diff
+
+# make sure that the resulting model has no more than 256 different values (i.e. quantized)
+$MRT_TOOLS/check-model-unique-vals.py train/model.npz -b 8 --with_bias
+
+# Exit with success code
+exit 0
diff --git a/tests/training/features/quantized-model/test_quantmodel_with_optimization.sh b/tests/training/features/quantized-model/test_quantmodel_with_optimization.sh
new file mode 100644
index 0000000..fe7993e
--- /dev/null
+++ b/tests/training/features/quantized-model/test_quantmodel_with_optimization.sh
@@ -0,0 +1,36 @@
+#!/bin/bash -x
+
+#####################################################################
+# SUMMARY: Train a quantized marian model with a scale optimization
+# AUTHOR: afaji
+#####################################################################
+
+# Exit on error
+set -e
+
+PREFIX=quantized-opt
+
+# Remove old artifacts and create working directory
+rm -rf train $PREFIX.{log,out,diff}
+mkdir -p train
+
+# Train an 8-bits model
+$MRT_MARIAN/marian \
+    --no-shuffle --seed 1111 --dim-emb 32 --dim-rnn 64 --mini-batch 32 --maxi-batch 1 --maxi-batch-sort none --learn-rate 0.1 --optimizer sgd \
+    -m train/model.npz -t $MRT_DATA/europarl.de-en/corpus.bpe.{en,de} -v train/vocab.en.yml train/vocab.de.yml \
+    --cost-type cross-entropy --sync-sgd --after-batches 100 --disp-freq 10 --quantize-bits 8 --quantize-optimization-steps 3 \
+    --log $PREFIX.log
+
+# Check if files exist
+test -e train/model.npz
+test -e $PREFIX.log
+
+# Compare the current output with the expected output
+cat $PREFIX.log | $MRT_TOOLS/extract-costs.sh > $PREFIX.out
+$MRT_TOOLS/diff-nums.py $PREFIX.out $PREFIX.expected -o $PREFIX.diff
+
+# make sure that the resulting model has no more than 256 different values (i.e. quantized)
+$MRT_TOOLS/check-model-unique-vals.py train/model.npz -b 8
+
+# Exit with success code
+exit 0
diff --git a/tools/check-model-unique-vals.py b/tools/check-model-unique-vals.py
new file mode 100755
index 0000000..59516f2
--- /dev/null
+++ b/tools/check-model-unique-vals.py
@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import os
+import sys
+import argparse
+import re
+
+import numpy as np
+
+def main():
+    exit_code = 0
+    args = parse_user_args()
+
+    with np.load(args.file) as data:
+      for key in data:
+        # skip special:model.yml
+        if "special" in key:
+          continue
+       
+        # if one of the dimension is 1, then it is a bias
+        # skip if it is bias and bias is not included
+        smallest_dim = sorted(data[key].shape)[0]
+        if(smallest_dim == 1 and not args.with_bias):
+          continue
+
+        if (np.unique(data[key]).size > 2**args.bits):
+          message("Tensor {} has more than {} unique values".format( \
+                   key, \
+                   2**args.bits), args)
+          exit_code = 1
+        if (args.print_centers):
+          message("Tensor {} unique centers: {}".format(key, np.unique(data[key])), args)    
+    return exit_code
+
+
+def message(text, args):
+    if not text.endswith("\n"):
+        text += "\n"
+    args.output.write(text)
+    if not args.quiet \
+            and args.output is not sys.stdout \
+            and args.output is not sys.stderr:
+        sys.stderr.write(text)
+
+
+def parse_user_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("file", type=str)
+    parser.add_argument("-o", "--output", type=argparse.FileType('w'), metavar="FILE", default=sys.stdout)
+    parser.add_argument("--print_centers", action="store_true")
+    parser.add_argument("-b", "--bits", type=int)
+    parser.add_argument("--with_bias", action="store_true") 
+    parser.add_argument("-q", "--quiet", action="store_true")
+    return parser.parse_args()
+
+if __name__ == '__main__':
+    code = main()
+    exit(code)