fix some bugs

Author: chenxi

.gitignore

@@ -3,4 +3,6 @@ __pycache__
 *.jpg
 *.png
 log
+pytorch_playground.egg-info
+script/val224_compressed.pkl
 

README.md

@@ -24,27 +24,21 @@ Also, if want to train the MLP model on mnist, simply run `python mnist/train.py
 
 
 # Install
-- pytorch (>=0.1.11) and torchvision from [official website](http://pytorch.org/), for example, cuda8.0 for python3.5
-    - `pip install http://download.pytorch.org/whl/cu80/torch-0.1.12.post2-cp35-cp35m-linux_x86_64.whl`
-    - `pip install torchvision`
-- tqdm
-    - `pip install tqdm`
-- OpenCV
-    - `conda install -c menpo opencv3`
-- Setting PYTHONPATH
-    - `export PYTHONPATH=/path/to/pytorch-playground:$PYTHONPATH`
+```
+python3 setup.py develop --user
+```
 
 # ImageNet dataset
 We provide precomputed imagenet validation dataset with 224x224x3 size. We first resize the shorter size of image to 256, then we crop 224x224 image in the center. Then we encode the cropped images to jpg string and dump to pickle. 
 - `cd script`
-- Download the [val224_compressed.pkl](https://drive.google.com/file/d/1U8ir2fOR4Sir3FCj9b7FQRPSVsycTfVc/view?usp=sharing) 
-- `python convert.py`
+- Download the `val224_compressed.pkl` ([Tsinghua](http://ml.cs.tsinghua.edu.cn/~chenxi/dataset/val224_compressed.pkl) /  [Google Drive](https://drive.google.com/file/d/1U8ir2fOR4Sir3FCj9b7FQRPSVsycTfVc/view?usp=sharing))
+- `python convert.py` (needs 48G memory, thanks [@jnorwood](https://github.com/aaron-xichen/pytorch-playground/issues/18) )
 
 
 # Quantization
 We also provide a simple demo to quantize these models to specified bit-width with several methods, including linear method, minmax method and non-linear method.
 
-`python quantize.py --type cifar10 --quant_method linear --param_bits 8 --fwd_bits 8 --bn_bits 8 --ngpu 1`
+`quantize --type cifar10 --quant_method linear --param_bits 8 --fwd_bits 8 --bn_bits 8 --ngpu 1`
 
 ## Top1 Accuracy
 We evaluate the performance of popular dataset and models with linear quantized method. The bit-width of running mean and running variance in BN are 10 bits for all results. (except for 32-float)

imagenet/dataset.py

@@ -2,6 +2,7 @@
 import os
 import os.path
 import numpy as np
+import joblib
 
 
 def get(batch_size, data_root='/tmp/public_dataset/pytorch', train=False, val=True, **kwargs):
@@ -26,7 +27,7 @@ def __init__(self, root, batch_size, train=False, input_size=224, **kwargs):
             pkl_file = os.path.join(root, 'train{}.pkl'.format(input_size))
         else:
             pkl_file = os.path.join(root, 'val{}.pkl'.format(input_size))
-        self.data_dict = misc.load_pickle(pkl_file)
+        self.data_dict = joblib.load(pkl_file)
 
         self.batch_size = batch_size
         self.idx = 0

imagenet/inception.py

@@ -60,7 +60,7 @@ def __init__(self, num_classes=1000, aux_logits=True, transform_input=False):
                 stddev = m.stddev if hasattr(m, 'stddev') else 0.1
                 X = stats.truncnorm(-2, 2, scale=stddev)
                 values = torch.Tensor(X.rvs(m.weight.data.numel()))
-                m.weight.data.copy_(values)
+                m.weight.data.copy_(values.reshape(m.weight.shape))
             elif isinstance(m, nn.BatchNorm2d):
                 m.weight.data.fill_(1)
                 m.bias.data.zero_()

quantize.py

@@ -5,93 +5,95 @@
 cudnn.benchmark =True
 from collections import OrderedDict
 
-parser = argparse.ArgumentParser(description='PyTorch SVHN Example')
-parser.add_argument('--type', default='cifar10', help='|'.join(selector.known_models))
-parser.add_argument('--quant_method', default='linear', help='linear|minmax|log|tanh')
-parser.add_argument('--batch_size', type=int, default=100, help='input batch size for training (default: 64)')
-parser.add_argument('--gpu', default=None, help='index of gpus to use')
-parser.add_argument('--ngpu', type=int, default=8, help='number of gpus to use')
-parser.add_argument('--seed', type=int, default=117, help='random seed (default: 1)')
-parser.add_argument('--model_root', default='~/.torch/models/', help='folder to save the model')
-parser.add_argument('--data_root', default='/tmp/public_dataset/pytorch/', help='folder to save the model')
-parser.add_argument('--logdir', default='log/default', help='folder to save to the log')
+def main():
+    parser = argparse.ArgumentParser(description='PyTorch SVHN Example')
+    parser.add_argument('--type', default='cifar10', help='|'.join(selector.known_models))
+    parser.add_argument('--quant_method', default='linear', help='linear|minmax|log|tanh')
+    parser.add_argument('--batch_size', type=int, default=100, help='input batch size for training (default: 64)')
+    parser.add_argument('--gpu', default=None, help='index of gpus to use')
+    parser.add_argument('--ngpu', type=int, default=8, help='number of gpus to use')
+    parser.add_argument('--seed', type=int, default=117, help='random seed (default: 1)')
+    parser.add_argument('--model_root', default='~/.torch/models/', help='folder to save the model')
+    parser.add_argument('--data_root', default='/data/public_dataset/pytorch/', help='folder to save the model')
+    parser.add_argument('--logdir', default='log/default', help='folder to save to the log')
 
-parser.add_argument('--input_size', type=int, default=224, help='input size of image')
-parser.add_argument('--n_sample', type=int, default=20, help='number of samples to infer the scaling factor')
-parser.add_argument('--param_bits', type=int, default=8, help='bit-width for parameters')
-parser.add_argument('--bn_bits', type=int, default=32, help='bit-width for running mean and std')
-parser.add_argument('--fwd_bits', type=int, default=8, help='bit-width for layer output')
-parser.add_argument('--overflow_rate', type=float, default=0.0, help='overflow rate')
-args = parser.parse_args()
+    parser.add_argument('--input_size', type=int, default=224, help='input size of image')
+    parser.add_argument('--n_sample', type=int, default=20, help='number of samples to infer the scaling factor')
+    parser.add_argument('--param_bits', type=int, default=8, help='bit-width for parameters')
+    parser.add_argument('--bn_bits', type=int, default=32, help='bit-width for running mean and std')
+    parser.add_argument('--fwd_bits', type=int, default=8, help='bit-width for layer output')
+    parser.add_argument('--overflow_rate', type=float, default=0.0, help='overflow rate')
+    args = parser.parse_args()
 
-args.gpu = misc.auto_select_gpu(utility_bound=0, num_gpu=args.ngpu, selected_gpus=args.gpu)
-args.ngpu = len(args.gpu)
-misc.ensure_dir(args.logdir)
-args.model_root = misc.expand_user(args.model_root)
-args.data_root = misc.expand_user(args.data_root)
-args.input_size = 299 if 'inception' in args.type else args.input_size
-assert args.quant_method in ['linear', 'minmax', 'log', 'tanh']
-print("=================FLAGS==================")
-for k, v in args.__dict__.items():
-    print('{}: {}'.format(k, v))
-print("========================================")
+    args.gpu = misc.auto_select_gpu(utility_bound=0, num_gpu=args.ngpu, selected_gpus=args.gpu)
+    args.ngpu = len(args.gpu)
+    misc.ensure_dir(args.logdir)
+    args.model_root = misc.expand_user(args.model_root)
+    args.data_root = misc.expand_user(args.data_root)
+    args.input_size = 299 if 'inception' in args.type else args.input_size
+    assert args.quant_method in ['linear', 'minmax', 'log', 'tanh']
+    print("=================FLAGS==================")
+    for k, v in args.__dict__.items():
+        print('{}: {}'.format(k, v))
+    print("========================================")
 
-assert torch.cuda.is_available(), 'no cuda'
-torch.manual_seed(args.seed)
-torch.cuda.manual_seed(args.seed)
+    assert torch.cuda.is_available(), 'no cuda'
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed(args.seed)
 
-# load model and dataset fetcher
-model_raw, ds_fetcher, is_imagenet = selector.select(args.type, model_root=args.model_root)
-args.ngpu = args.ngpu if is_imagenet else 1
+    # load model and dataset fetcher
+    model_raw, ds_fetcher, is_imagenet = selector.select(args.type, model_root=args.model_root)
+    args.ngpu = args.ngpu if is_imagenet else 1
 
-# quantize parameters
-if args.param_bits < 32:
-    state_dict = model_raw.state_dict()
-    state_dict_quant = OrderedDict()
-    sf_dict = OrderedDict()
-    for k, v in state_dict.items():
-        if 'running' in k:
-            if args.bn_bits >=32:
-                print("Ignoring {}".format(k))
-                state_dict_quant[k] = v
-                continue
+    # quantize parameters
+    if args.param_bits < 32:
+        state_dict = model_raw.state_dict()
+        state_dict_quant = OrderedDict()
+        sf_dict = OrderedDict()
+        for k, v in state_dict.items():
+            if 'running' in k:
+                if args.bn_bits >=32:
+                    print("Ignoring {}".format(k))
+                    state_dict_quant[k] = v
+                    continue
+                else:
+                    bits = args.bn_bits
             else:
-                bits = args.bn_bits
-        else:
-            bits = args.param_bits
+                bits = args.param_bits
 
-        if args.quant_method == 'linear':
-            sf = bits - 1. - quant.compute_integral_part(v, overflow_rate=args.overflow_rate)
-            v_quant  = quant.linear_quantize(v, sf, bits=bits)
-        elif args.quant_method == 'log':
-            v_quant = quant.log_minmax_quantize(v, bits=bits)
-        elif args.quant_method == 'minmax':
-            v_quant = quant.min_max_quantize(v, bits=bits)
-        else:
-            v_quant = quant.tanh_quantize(v, bits=bits)
-        state_dict_quant[k] = v_quant
-        print(k, bits)
-    model_raw.load_state_dict(state_dict_quant)
-
-# quantize forward activation
-if args.fwd_bits < 32:
-    model_raw = quant.duplicate_model_with_quant(model_raw, bits=args.fwd_bits, overflow_rate=args.overflow_rate,
-                                                 counter=args.n_sample, type=args.quant_method)
-    print(model_raw)
-    val_ds_tmp = ds_fetcher(10, data_root=args.data_root, train=False, input_size=args.input_size)
-    misc.eval_model(model_raw, val_ds_tmp, ngpu=1, n_sample=args.n_sample, is_imagenet=is_imagenet)
+            if args.quant_method == 'linear':
+                sf = bits - 1. - quant.compute_integral_part(v, overflow_rate=args.overflow_rate)
+                v_quant  = quant.linear_quantize(v, sf, bits=bits)
+            elif args.quant_method == 'log':
+                v_quant = quant.log_minmax_quantize(v, bits=bits)
+            elif args.quant_method == 'minmax':
+                v_quant = quant.min_max_quantize(v, bits=bits)
+            else:
+                v_quant = quant.tanh_quantize(v, bits=bits)
+            state_dict_quant[k] = v_quant
+            print(k, bits)
+        model_raw.load_state_dict(state_dict_quant)
 
-# eval model
-val_ds = ds_fetcher(args.batch_size, data_root=args.data_root, train=False, input_size=args.input_size)
-acc1, acc5 = misc.eval_model(model_raw, val_ds, ngpu=args.ngpu, is_imagenet=is_imagenet)
+    # quantize forward activation
+    if args.fwd_bits < 32:
+        model_raw = quant.duplicate_model_with_quant(model_raw, bits=args.fwd_bits, overflow_rate=args.overflow_rate,
+                                                     counter=args.n_sample, type=args.quant_method)
+        print(model_raw)
+        val_ds_tmp = ds_fetcher(10, data_root=args.data_root, train=False, input_size=args.input_size)
+        misc.eval_model(model_raw, val_ds_tmp, ngpu=1, n_sample=args.n_sample, is_imagenet=is_imagenet)
 
-# print sf
-print(model_raw)
-res_str = "type={}, quant_method={}, param_bits={}, bn_bits={}, fwd_bits={}, overflow_rate={}, acc1={:.4f}, acc5={:.4f}".format(
-    args.type, args.quant_method, args.param_bits, args.bn_bits, args.fwd_bits, args.overflow_rate, acc1, acc5)
-print(res_str)
-with open('acc1_acc5.txt', 'a') as f:
-    f.write(res_str + '\n')
+    # eval model
+    val_ds = ds_fetcher(args.batch_size, data_root=args.data_root, train=False, input_size=args.input_size)
+    acc1, acc5 = misc.eval_model(model_raw, val_ds, ngpu=args.ngpu, is_imagenet=is_imagenet)
 
+    # print sf
+    print(model_raw)
+    res_str = "type={}, quant_method={}, param_bits={}, bn_bits={}, fwd_bits={}, overflow_rate={}, acc1={:.4f}, acc5={:.4f}".format(
+        args.type, args.quant_method, args.param_bits, args.bn_bits, args.fwd_bits, args.overflow_rate, acc1, acc5)
+    print(res_str)
+    with open('acc1_acc5.txt', 'a') as f:
+        f.write(res_str + '\n')
 
 
+if __name__ == '__main__':
+    main()

script/convert.py

@@ -4,41 +4,45 @@
 from utee import misc
 import argparse
 import cv2
+import joblib
 
-imagenet_urls = [
-   'http://ml.cs.tsinghua.edu.cn/~chenxi/dataset/val224_compressed.pkl'
-]
 parser = argparse.ArgumentParser(description='Extract the ILSVRC2012 val dataset')
 parser.add_argument('--in_file', default='val224_compressed.pkl', help='input file path')
-parser.add_argument('--out_root', default='/tmp/public_dataset/pytorch/imagenet-data/', help='output file path')
+parser.add_argument('--out_root', default='/data/public_dataset/pytorch/imagenet-data/', help='output file path')
 args = parser.parse_args()
 
 d = misc.load_pickle(args.in_file)
 assert len(d['data']) == 50000, len(d['data'])
 assert len(d['target']) == 50000, len(d['target'])
 
-data224 = []
+
 data299 = []
 for img, target in tqdm.tqdm(zip(d['data'], d['target']), total=50000):
     img224 = misc.str2img(img)
     img299 = cv2.resize(img224, (299, 299))
-    data224.append(img224)
     data299.append(img299)
-data_dict224 = dict(
-    data = np.array(data224).transpose(0, 3, 1, 2),
-    target = d['target']
-)
+
 data_dict299 = dict(
     data = np.array(data299).transpose(0, 3, 1, 2),
     target = d['target']
 )
-
 if not os.path.exists(args.out_root):
     os.makedirs(args.out_root)
-misc.dump_pickle(data_dict224, os.path.join(args.out_root, 'val224.pkl'))
-misc.dump_pickle(data_dict299, os.path.join(args.out_root, 'val299.pkl'))
+joblib.dump(data_dict299, os.path.join(args.out_root, 'val299.pkl'))
 
+data299.clear()
+data_dict299.clear()
 
+data224 = []
+for img, target in tqdm.tqdm(zip(d['data'], d['target']), total=50000):
+    img224 = misc.str2img(img)
+    data224.append(img224)
+
+data_dict224 = dict(
+    data = np.array(data224).transpose(0, 3, 1, 2),
+    target = d['target']
+)
+joblib.dump(data_dict224, os.path.join(args.out_root, 'val224.pkl'))
 
 
 

setup.py

@@ -0,0 +1,19 @@
+from setuptools import setup, find_packages
+
+with open("requirements.txt") as requirements_file:
+    REQUIREMENTS = requirements_file.readlines()
+
+setup(
+    name="pytorch-playground",
+    version="1.0.0",
+    author='Aaron Chen',
+    author_email='[email protected]',
+    packages=find_packages(),
+    entry_points = {
+        'console_scripts': [
+            'quantize=quantize:main',
+        ]
+    },
+    install_requires=REQUIREMENTS,
+
+)

svhn/dataset.py

@@ -10,7 +10,7 @@ def get(batch_size, data_root='/tmp/public_dataset/pytorch', train=True, val=Tru
     print("Building SVHN data loader with {} workers".format(num_workers))
 
     def target_transform(target):
-        return int(target[0]) - 1
+        return int(target) - 1
 
     ds = []
     if train:

utee/misc.py

@@ -222,6 +222,7 @@ def load_state_dict(model, model_urls, model_root):
         own_state[name].copy_(param)
 
     missing = set(own_state.keys()) - set(state_dict.keys())
-    if len(missing) > 0:
-        raise KeyError('missing keys in state_dict: "{}"'.format(missing))
+    no_use = set(state_dict.keys()) - set(own_state.keys())
+    if len(no_use) > 0:
+        raise KeyError('some keys are not used: "{}"'.format(no_use))
 

utee/quant.py

@@ -11,7 +11,7 @@ def compute_integral_part(input, overflow_rate):
     split_idx = int(overflow_rate * len(sorted_value))
     v = sorted_value[split_idx]
     if isinstance(v, Variable):
-        v = v.data.cpu().numpy()[0]
+        v = float(v.data.cpu())
     sf = math.ceil(math.log2(v+1e-12))
     return sf
 
@@ -35,7 +35,7 @@ def log_minmax_quantize(input, bits):
 
     s = torch.sign(input)
     input0 = torch.log(torch.abs(input) + 1e-20)
-    v = min_max_quantize(input0, bits)
+    v = min_max_quantize(input0, bits-1)
     v = torch.exp(v) * s
     return v
 
@@ -46,7 +46,7 @@ def log_linear_quantize(input, sf, bits):
 
     s = torch.sign(input)
     input0 = torch.log(torch.abs(input) + 1e-20)
-    v = linear_quantize(input0, sf, bits)
+    v = linear_quantize(input0, sf, bits-1)
     v = torch.exp(v) * s
     return v