├── .gitignore ├── LICENSE ├── README.md ├── cifar ├── __init__.py ├── dataset.py ├── model.py └── train.py ├── imagenet ├── __init__.py ├── alexnet.py ├── dataset.py ├── inception.py ├── resnet.py ├── squeezenet.py └── vgg.py ├── mnist ├── __init__.py ├── dataset.py ├── model.py └── train.py ├── quantize.py ├── requirements.txt ├── roadmap_zh.md ├── script └── convert.py ├── setup.py ├── stl10 ├── __init__.py ├── dataset.py ├── model.py └── train.py ├── svhn ├── __init__.py ├── dataset.py ├── model.py └── train.py └── utee ├── __init__.py ├── misc.py ├── quant.py └── selector.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | *.jpg 3 | *.png 4 | acc1_acc5.txt 5 | log 6 | pytorch_playground.egg-info 7 | script/val224_compressed.pkl 8 | 9 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 Aaron Chen 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | This is a playground for pytorch beginners, which contains predefined models on popular dataset. Currently we support 2 | - mnist, svhn 3 | - cifar10, cifar100 4 | - stl10 5 | - alexnet 6 | - vgg16, vgg16_bn, vgg19, vgg19_bn 7 | - resnet18, resnet34, resnet50, resnet101, resnet152 8 | - squeezenet_v0, squeezenet_v1 9 | - inception_v3 10 | 11 | Here is an example for MNIST dataset. This will download the dataset and pre-trained model automatically. 12 | ``` 13 | import torch 14 | from torch.autograd import Variable 15 | from utee import selector 16 | model_raw, ds_fetcher, is_imagenet = selector.select('mnist') 17 | ds_val = ds_fetcher(batch_size=10, train=False, val=True) 18 | for idx, (data, target) in enumerate(ds_val): 19 | data = Variable(torch.FloatTensor(data)).cuda() 20 | output = model_raw(data) 21 | ``` 22 | 23 | Also, if want to train the MLP model on mnist, simply run `python mnist/train.py` 24 | 25 | 26 | # Install 27 | ``` 28 | python3 setup.py develop --user 29 | ``` 30 | 31 | # ImageNet dataset 32 | 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. 33 | - `cd script` 34 | - 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)) 35 | - `python convert.py` (needs 48G memory, thanks [@jnorwood](https://github.com/aaron-xichen/pytorch-playground/issues/18) ) 36 | 37 | 38 | # Quantization 39 | 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. 40 | 41 | `quantize --type cifar10 --quant_method linear --param_bits 8 --fwd_bits 8 --bn_bits 8 --ngpu 1` 42 | 43 | ## Top1 Accuracy 44 | 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) 45 | 46 | 47 | |Model|32-float |12-bit |10-bit |8-bit |6-bit | 48 | |:----|:--------:|:------:|:-----:|:-----:|:-----:| 49 | |[MNIST](http://ml.cs.tsinghua.edu.cn/~chenxi/pytorch-models/mnist-b07bb66b.pth)|98.42|98.43|98.44|98.44|98.32| 50 | |[SVHN](http://ml.cs.tsinghua.edu.cn/~chenxi/pytorch-models/svhn-f564f3d8.pth)|96.03|96.03|96.04|96.02|95.46| 51 | |[CIFAR10](http://ml.cs.tsinghua.edu.cn/~chenxi/pytorch-models/cifar10-d875770b.pth)|93.78|93.79|93.80|93.58|90.86| 52 | |[CIFAR100](http://ml.cs.tsinghua.edu.cn/~chenxi/pytorch-models/cifar100-3a55a987.pth)|74.27|74.21|74.19|73.70|66.32| 53 | |[STL10](http://ml.cs.tsinghua.edu.cn/~chenxi/pytorch-models/stl10-866321e9.pth)|77.59|77.65|77.70|77.59|73.40| 54 | |[AlexNet](https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth)|55.70/78.42|55.66/78.41|55.54/78.39|54.17/77.29|18.19/36.25| 55 | |[VGG16](https://download.pytorch.org/models/vgg16-397923af.pth)|70.44/89.43|70.45/89.43|70.44/89.33|69.99/89.17|53.33/76.32| 56 | |[VGG19](https://download.pytorch.org/models/vgg19-dcbb9e9d.pth)|71.36/89.94|71.35/89.93|71.34/89.88|70.88/89.62|56.00/78.62| 57 | |[ResNet18](https://download.pytorch.org/models/resnet18-5c106cde.pth)|68.63/88.31|68.62/88.33|68.49/88.25|66.80/87.20|19.14/36.49| 58 | |[ResNet34](https://download.pytorch.org/models/resnet34-333f7ec4.pth)|72.50/90.86|72.46/90.82|72.45/90.85|71.47/90.00|32.25/55.71| 59 | |[ResNet50](https://download.pytorch.org/models/resnet50-19c8e357.pth)|74.98/92.17|74.94/92.12|74.91/92.09|72.54/90.44|2.43/5.36| 60 | |[ResNet101](https://download.pytorch.org/models/resnet101-5d3b4d8f.pth)|76.69/93.30|76.66/93.25|76.22/92.90|65.69/79.54|1.41/1.18| 61 | |[ResNet152](https://download.pytorch.org/models/resnet152-b121ed2d.pth)|77.55/93.59|77.51/93.62|77.40/93.54|74.95/92.46|9.29/16.75| 62 | |[SqueezeNetV0](https://download.pytorch.org/models/squeezenet1_0-a815701f.pth)|56.73/79.39|56.75/79.40|56.70/79.27|53.93/77.04|14.21/29.74| 63 | |[SqueezeNetV1](https://download.pytorch.org/models/squeezenet1_1-f364aa15.pth)|56.52/79.13|56.52/79.15|56.24/79.03|54.56/77.33|17.10/32.46| 64 | |[InceptionV3](https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth)|76.41/92.78|76.43/92.71|76.44/92.73|73.67/91.34|1.50/4.82| 65 | 66 | **Note: ImageNet 32-float models are directly from torchvision** 67 | 68 | 69 | ## Selected Arguments 70 | Here we give an overview of selected arguments of `quantize.py` 71 | 72 | |Flag |Default value|Description & Options| 73 | |:-----------------------------|:-----------------------:|:--------------------------------| 74 | |type|cifar10|mnist,svhn,cifar10,cifar100,stl10,alexnet,vgg16,vgg16_bn,vgg19,vgg19_bn,resent18,resent34,resnet50,resnet101,resnet152,squeezenet_v0,squeezenet_v1,inception_v3| 75 | |quant_method|linear|quantization method:linear,minmax,log,tanh| 76 | |param_bits|8|bit-width of weights and bias| 77 | |fwd_bits|8|bit-width of activation| 78 | |bn_bits|32|bit-width of running mean and running vairance| 79 | |overflow_rate|0.0|overflow rate threshold for linear quantization method| 80 | |n_samples|20|number of samples to make statistics for activation| 81 | -------------------------------------------------------------------------------- /cifar/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aaron-xichen/pytorch-playground/5d476295f1c1564c8074eb885bb6138d6bd09889/cifar/__init__.py -------------------------------------------------------------------------------- /cifar/dataset.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torchvision import datasets, transforms 3 | from torch.utils.data import DataLoader 4 | import os 5 | 6 | def get10(batch_size, data_root='/tmp/public_dataset/pytorch', train=True, val=True, **kwargs): 7 | data_root = os.path.expanduser(os.path.join(data_root, 'cifar10-data')) 8 | num_workers = kwargs.setdefault('num_workers', 1) 9 | kwargs.pop('input_size', None) 10 | print("Building CIFAR-10 data loader with {} workers".format(num_workers)) 11 | ds = [] 12 | if train: 13 | train_loader = torch.utils.data.DataLoader( 14 | datasets.CIFAR10( 15 | root=data_root, train=True, download=True, 16 | transform=transforms.Compose([ 17 | transforms.Pad(4), 18 | transforms.RandomCrop(32), 19 | transforms.RandomHorizontalFlip(), 20 | transforms.ToTensor(), 21 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 22 | ])), 23 | batch_size=batch_size, shuffle=True, **kwargs) 24 | ds.append(train_loader) 25 | if val: 26 | test_loader = torch.utils.data.DataLoader( 27 | datasets.CIFAR10( 28 | root=data_root, train=False, download=True, 29 | transform=transforms.Compose([ 30 | transforms.ToTensor(), 31 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 32 | ])), 33 | batch_size=batch_size, shuffle=False, **kwargs) 34 | ds.append(test_loader) 35 | ds = ds[0] if len(ds) == 1 else ds 36 | return ds 37 | 38 | def get100(batch_size, data_root='/tmp/public_dataset/pytorch', train=True, val=True, **kwargs): 39 | data_root = os.path.expanduser(os.path.join(data_root, 'cifar100-data')) 40 | num_workers = kwargs.setdefault('num_workers', 1) 41 | kwargs.pop('input_size', None) 42 | print("Building CIFAR-100 data loader with {} workers".format(num_workers)) 43 | ds = [] 44 | if train: 45 | train_loader = torch.utils.data.DataLoader( 46 | datasets.CIFAR100( 47 | root=data_root, train=True, download=True, 48 | transform=transforms.Compose([ 49 | transforms.Pad(4), 50 | transforms.RandomCrop(32), 51 | transforms.RandomHorizontalFlip(), 52 | transforms.ToTensor(), 53 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 54 | ])), 55 | batch_size=batch_size, shuffle=True, **kwargs) 56 | ds.append(train_loader) 57 | 58 | if val: 59 | test_loader = torch.utils.data.DataLoader( 60 | datasets.CIFAR100( 61 | root=data_root, train=False, download=True, 62 | transform=transforms.Compose([ 63 | transforms.ToTensor(), 64 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 65 | ])), 66 | batch_size=batch_size, shuffle=False, **kwargs) 67 | ds.append(test_loader) 68 | ds = ds[0] if len(ds) == 1 else ds 69 | return ds 70 | 71 | -------------------------------------------------------------------------------- /cifar/model.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.utils.model_zoo as model_zoo 3 | from IPython import embed 4 | from collections import OrderedDict 5 | 6 | from utee import misc 7 | print = misc.logger.info 8 | 9 | model_urls = { 10 | 'cifar10': 'http://ml.cs.tsinghua.edu.cn/~chenxi/pytorch-models/cifar10-d875770b.pth', 11 | 'cifar100': 'http://ml.cs.tsinghua.edu.cn/~chenxi/pytorch-models/cifar100-3a55a987.pth', 12 | } 13 | 14 | class CIFAR(nn.Module): 15 | def __init__(self, features, n_channel, num_classes): 16 | super(CIFAR, self).__init__() 17 | assert isinstance(features, nn.Sequential), type(features) 18 | self.features = features 19 | self.classifier = nn.Sequential( 20 | nn.Linear(n_channel, num_classes) 21 | ) 22 | print(self.features) 23 | print(self.classifier) 24 | 25 | def forward(self, x): 26 | x = self.features(x) 27 | x = x.view(x.size(0), -1) 28 | x = self.classifier(x) 29 | return x 30 | 31 | def make_layers(cfg, batch_norm=False): 32 | layers = [] 33 | in_channels = 3 34 | for i, v in enumerate(cfg): 35 | if v == 'M': 36 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 37 | else: 38 | padding = v[1] if isinstance(v, tuple) else 1 39 | out_channels = v[0] if isinstance(v, tuple) else v 40 | conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=padding) 41 | if batch_norm: 42 | layers += [conv2d, nn.BatchNorm2d(out_channels, affine=False), nn.ReLU()] 43 | else: 44 | layers += [conv2d, nn.ReLU()] 45 | in_channels = out_channels 46 | return nn.Sequential(*layers) 47 | 48 | def cifar10(n_channel, pretrained=None): 49 | cfg = [n_channel, n_channel, 'M', 2*n_channel, 2*n_channel, 'M', 4*n_channel, 4*n_channel, 'M', (8*n_channel, 0), 'M'] 50 | layers = make_layers(cfg, batch_norm=True) 51 | model = CIFAR(layers, n_channel=8*n_channel, num_classes=10) 52 | if pretrained is not None: 53 | m = model_zoo.load_url(model_urls['cifar10']) 54 | state_dict = m.state_dict() if isinstance(m, nn.Module) else m 55 | assert isinstance(state_dict, (dict, OrderedDict)), type(state_dict) 56 | model.load_state_dict(state_dict) 57 | return model 58 | 59 | def cifar100(n_channel, pretrained=None): 60 | cfg = [n_channel, n_channel, 'M', 2*n_channel, 2*n_channel, 'M', 4*n_channel, 4*n_channel, 'M', (8*n_channel, 0), 'M'] 61 | layers = make_layers(cfg, batch_norm=True) 62 | model = CIFAR(layers, n_channel=8*n_channel, num_classes=100) 63 | if pretrained is not None: 64 | m = model_zoo.load_url(model_urls['cifar100']) 65 | state_dict = m.state_dict() if isinstance(m, nn.Module) else m 66 | assert isinstance(state_dict, (dict, OrderedDict)), type(state_dict) 67 | model.load_state_dict(state_dict) 68 | return model 69 | 70 | if __name__ == '__main__': 71 | model = cifar10(128, pretrained='log/cifar10/best-135.pth') 72 | embed() 73 | 74 | -------------------------------------------------------------------------------- /cifar/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import time 4 | 5 | from utee import misc 6 | import torch 7 | import torch.nn.functional as F 8 | import torch.optim as optim 9 | from torch.autograd import Variable 10 | 11 | import dataset 12 | import model 13 | 14 | from IPython import embed 15 | 16 | parser = argparse.ArgumentParser(description='PyTorch CIFAR-X Example') 17 | parser.add_argument('--type', default='cifar10', help='cifar10|cifar100') 18 | parser.add_argument('--channel', type=int, default=128, help='first conv channel (default: 32)') 19 | parser.add_argument('--wd', type=float, default=0.00, help='weight decay') 20 | parser.add_argument('--batch_size', type=int, default=200, help='input batch size for training (default: 64)') 21 | parser.add_argument('--epochs', type=int, default=150, help='number of epochs to train (default: 10)') 22 | parser.add_argument('--lr', type=float, default=0.001, help='learning rate (default: 1e-3)') 23 | parser.add_argument('--gpu', default=None, help='index of gpus to use') 24 | parser.add_argument('--ngpu', type=int, default=2, help='number of gpus to use') 25 | parser.add_argument('--seed', type=int, default=117, help='random seed (default: 1)') 26 | parser.add_argument('--log_interval', type=int, default=100, help='how many batches to wait before logging training status') 27 | parser.add_argument('--test_interval', type=int, default=5, help='how many epochs to wait before another test') 28 | parser.add_argument('--logdir', default='log/default', help='folder to save to the log') 29 | parser.add_argument('--decreasing_lr', default='80,120', help='decreasing strategy') 30 | args = parser.parse_args() 31 | args.logdir = os.path.join(os.path.dirname(__file__), args.logdir) 32 | misc.logger.init(args.logdir, 'train_log') 33 | print = misc.logger.info 34 | 35 | # select gpu 36 | args.gpu = misc.auto_select_gpu(utility_bound=0, num_gpu=args.ngpu, selected_gpus=args.gpu) 37 | args.ngpu = len(args.gpu) 38 | 39 | # logger 40 | misc.ensure_dir(args.logdir) 41 | print("=================FLAGS==================") 42 | for k, v in args.__dict__.items(): 43 | print('{}: {}'.format(k, v)) 44 | print("========================================") 45 | 46 | # seed 47 | args.cuda = torch.cuda.is_available() 48 | torch.manual_seed(args.seed) 49 | if args.cuda: 50 | torch.cuda.manual_seed(args.seed) 51 | 52 | # data loader and model 53 | assert args.type in ['cifar10', 'cifar100'], args.type 54 | if args.type == 'cifar10': 55 | train_loader, test_loader = dataset.get10(batch_size=args.batch_size, num_workers=1) 56 | model = model.cifar10(n_channel=args.channel) 57 | else: 58 | train_loader, test_loader = dataset.get100(batch_size=args.batch_size, num_workers=1) 59 | model = model.cifar100(n_channel=args.channel) 60 | model = torch.nn.DataParallel(model, device_ids= range(args.ngpu)) 61 | if args.cuda: 62 | model.cuda() 63 | 64 | # optimizer 65 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd) 66 | decreasing_lr = list(map(int, args.decreasing_lr.split(','))) 67 | print('decreasing_lr: ' + str(decreasing_lr)) 68 | best_acc, old_file = 0, None 69 | t_begin = time.time() 70 | try: 71 | # ready to go 72 | for epoch in range(args.epochs): 73 | model.train() 74 | if epoch in decreasing_lr: 75 | optimizer.param_groups[0]['lr'] *= 0.1 76 | for batch_idx, (data, target) in enumerate(train_loader): 77 | indx_target = target.clone() 78 | if args.cuda: 79 | data, target = data.cuda(), target.cuda() 80 | data, target = Variable(data), Variable(target) 81 | 82 | optimizer.zero_grad() 83 | output = model(data) 84 | loss = F.cross_entropy(output, target) 85 | loss.backward() 86 | optimizer.step() 87 | 88 | if batch_idx % args.log_interval == 0 and batch_idx > 0: 89 | pred = output.data.max(1)[1] # get the index of the max log-probability 90 | correct = pred.cpu().eq(indx_target).sum() 91 | acc = correct * 1.0 / len(data) 92 | print('Train Epoch: {} [{}/{}] Loss: {:.6f} Acc: {:.4f} lr: {:.2e}'.format( 93 | epoch, batch_idx * len(data), len(train_loader.dataset), 94 | loss.data[0], acc, optimizer.param_groups[0]['lr'])) 95 | 96 | elapse_time = time.time() - t_begin 97 | speed_epoch = elapse_time / (epoch + 1) 98 | speed_batch = speed_epoch / len(train_loader) 99 | eta = speed_epoch * args.epochs - elapse_time 100 | print("Elapsed {:.2f}s, {:.2f} s/epoch, {:.2f} s/batch, ets {:.2f}s".format( 101 | elapse_time, speed_epoch, speed_batch, eta)) 102 | misc.model_snapshot(model, os.path.join(args.logdir, 'latest.pth')) 103 | 104 | if epoch % args.test_interval == 0: 105 | model.eval() 106 | test_loss = 0 107 | correct = 0 108 | for data, target in test_loader: 109 | indx_target = target.clone() 110 | if args.cuda: 111 | data, target = data.cuda(), target.cuda() 112 | data, target = Variable(data, volatile=True), Variable(target) 113 | output = model(data) 114 | test_loss += F.cross_entropy(output, target).data[0] 115 | pred = output.data.max(1)[1] # get the index of the max log-probability 116 | correct += pred.cpu().eq(indx_target).sum() 117 | 118 | test_loss = test_loss / len(test_loader) # average over number of mini-batch 119 | acc = 100. * correct / len(test_loader.dataset) 120 | print('\tTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format( 121 | test_loss, correct, len(test_loader.dataset), acc)) 122 | if acc > best_acc: 123 | new_file = os.path.join(args.logdir, 'best-{}.pth'.format(epoch)) 124 | misc.model_snapshot(model, new_file, old_file=old_file, verbose=True) 125 | best_acc = acc 126 | old_file = new_file 127 | except Exception as e: 128 | import traceback 129 | traceback.print_exc() 130 | finally: 131 | print("Total Elapse: {:.2f}, Best Result: {:.3f}%".format(time.time()-t_begin, best_acc)) 132 | 133 | 134 | -------------------------------------------------------------------------------- /imagenet/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aaron-xichen/pytorch-playground/5d476295f1c1564c8074eb885bb6138d6bd09889/imagenet/__init__.py -------------------------------------------------------------------------------- /imagenet/alexnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.utils.model_zoo as model_zoo 3 | 4 | 5 | __all__ = ['AlexNet', 'alexnet'] 6 | 7 | 8 | model_urls = { 9 | 'alexnet': 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth', 10 | } 11 | 12 | 13 | class AlexNet(nn.Module): 14 | 15 | def __init__(self, num_classes=1000): 16 | super(AlexNet, self).__init__() 17 | self.features = nn.Sequential( 18 | nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2), 19 | nn.ReLU(inplace=True), 20 | nn.MaxPool2d(kernel_size=3, stride=2), 21 | nn.Conv2d(64, 192, kernel_size=5, padding=2), 22 | nn.ReLU(inplace=True), 23 | nn.MaxPool2d(kernel_size=3, stride=2), 24 | nn.Conv2d(192, 384, kernel_size=3, padding=1), 25 | nn.ReLU(inplace=True), 26 | nn.Conv2d(384, 256, kernel_size=3, padding=1), 27 | nn.ReLU(inplace=True), 28 | nn.Conv2d(256, 256, kernel_size=3, padding=1), 29 | nn.ReLU(inplace=True), 30 | nn.MaxPool2d(kernel_size=3, stride=2), 31 | ) 32 | self.classifier = nn.Sequential( 33 | nn.Dropout(), 34 | nn.Linear(256 * 6 * 6, 4096), 35 | nn.ReLU(inplace=True), 36 | nn.Dropout(), 37 | nn.Linear(4096, 4096), 38 | nn.ReLU(inplace=True), 39 | nn.Linear(4096, num_classes), 40 | ) 41 | 42 | def forward(self, x): 43 | x = self.features(x) 44 | x = x.view(x.size(0), 256 * 6 * 6) 45 | x = self.classifier(x) 46 | return x 47 | 48 | 49 | def alexnet(pretrained=False, model_root=None, **kwargs): 50 | model = AlexNet(**kwargs) 51 | if pretrained: 52 | model.load_state_dict(model_zoo.load_url(model_urls['alexnet'], model_root)) 53 | return model 54 | -------------------------------------------------------------------------------- /imagenet/dataset.py: -------------------------------------------------------------------------------- 1 | from utee import misc 2 | import os 3 | import os.path 4 | import numpy as np 5 | import joblib 6 | 7 | 8 | def get(batch_size, data_root='/tmp/public_dataset/pytorch', train=False, val=True, **kwargs): 9 | data_root = os.path.expanduser(os.path.join(data_root, 'imagenet-data')) 10 | print("Building IMAGENET data loader, 50000 for train, 50000 for test") 11 | ds = [] 12 | assert train is not True, 'train not supported yet' 13 | if train: 14 | ds.append(IMAGENET(data_root, batch_size, True, **kwargs)) 15 | if val: 16 | ds.append(IMAGENET(data_root, batch_size, False, **kwargs)) 17 | ds = ds[0] if len(ds) == 1 else ds 18 | return ds 19 | 20 | class IMAGENET(object): 21 | def __init__(self, root, batch_size, train=False, input_size=224, **kwargs): 22 | self.mean = np.array([0.485, 0.456, 0.406]).reshape(1, 1, 1, 3) 23 | self.std = np.array([0.229, 0.224, 0.225]).reshape(1, 1, 1, 3) 24 | self.train = train 25 | 26 | if train: 27 | pkl_file = os.path.join(root, 'train{}.pkl'.format(input_size)) 28 | else: 29 | pkl_file = os.path.join(root, 'val{}.pkl'.format(input_size)) 30 | self.data_dict = joblib.load(pkl_file) 31 | 32 | self.batch_size = batch_size 33 | self.idx = 0 34 | 35 | @property 36 | def n_batch(self): 37 | return int(np.ceil(self.n_sample* 1.0 / self.batch_size)) 38 | 39 | @property 40 | def n_sample(self): 41 | return len(self.data_dict['data']) 42 | 43 | def __len__(self): 44 | return self.n_batch 45 | 46 | def __iter__(self): 47 | return self 48 | 49 | def __next__(self): 50 | if self.idx >= self.n_batch: 51 | self.idx = 0 52 | raise StopIteration 53 | else: 54 | img = self.data_dict['data'][self.idx*self.batch_size:(self.idx+1)*self.batch_size].astype('float32') 55 | target = self.data_dict['target'][self.idx*self.batch_size:(self.idx+1)*self.batch_size] 56 | self.idx += 1 57 | return img, target 58 | 59 | if __name__ == '__main__': 60 | train_ds, val_ds = get(200) 61 | 62 | 63 | -------------------------------------------------------------------------------- /imagenet/inception.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from utee import misc 5 | from collections import OrderedDict 6 | 7 | 8 | __all__ = ['Inception3', 'inception_v3'] 9 | 10 | 11 | model_urls = { 12 | 'inception_v3_google': 'https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth', 13 | } 14 | 15 | 16 | def inception_v3(pretrained=False, model_root=None, **kwargs): 17 | if pretrained: 18 | if 'transform_input' not in kwargs: 19 | kwargs['transform_input'] = True 20 | model = Inception3(**kwargs) 21 | misc.load_state_dict(model, model_urls['inception_v3_google'], model_root) 22 | return model 23 | 24 | return Inception3(**kwargs) 25 | 26 | 27 | class Inception3(nn.Module): 28 | 29 | def __init__(self, num_classes=1000, aux_logits=True, transform_input=False): 30 | super(Inception3, self).__init__() 31 | self.aux_logits = aux_logits 32 | self.transform_input = transform_input 33 | self.Conv2d_1a_3x3 = BasicConv2d(3, 32, kernel_size=3, stride=2) 34 | self.Conv2d_2a_3x3 = BasicConv2d(32, 32, kernel_size=3) 35 | self.Conv2d_2b_3x3 = BasicConv2d(32, 64, kernel_size=3, padding=1) 36 | self.Conv2d_3b_1x1 = BasicConv2d(64, 80, kernel_size=1) 37 | self.Conv2d_4a_3x3 = BasicConv2d(80, 192, kernel_size=3) 38 | self.Mixed_5b = InceptionA(192, pool_features=32) 39 | self.Mixed_5c = InceptionA(256, pool_features=64) 40 | self.Mixed_5d = InceptionA(288, pool_features=64) 41 | self.Mixed_6a = InceptionB(288) 42 | self.Mixed_6b = InceptionC(768, channels_7x7=128) 43 | self.Mixed_6c = InceptionC(768, channels_7x7=160) 44 | self.Mixed_6d = InceptionC(768, channels_7x7=160) 45 | self.Mixed_6e = InceptionC(768, channels_7x7=192) 46 | if aux_logits: 47 | self.AuxLogits = InceptionAux(768, num_classes) 48 | self.Mixed_7a = InceptionD(768) 49 | self.Mixed_7b = InceptionE(1280) 50 | self.Mixed_7c = InceptionE(2048) 51 | self.group1 = nn.Sequential( 52 | OrderedDict([ 53 | ('fc', nn.Linear(2048, num_classes)) 54 | ]) 55 | ) 56 | 57 | for m in self.modules(): 58 | if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): 59 | import scipy.stats as stats 60 | stddev = m.stddev if hasattr(m, 'stddev') else 0.1 61 | X = stats.truncnorm(-2, 2, scale=stddev) 62 | values = torch.Tensor(X.rvs(m.weight.data.numel())) 63 | m.weight.data.copy_(values.reshape(m.weight.shape)) 64 | elif isinstance(m, nn.BatchNorm2d): 65 | m.weight.data.fill_(1) 66 | m.bias.data.zero_() 67 | 68 | def forward(self, x): 69 | if self.transform_input: 70 | x = x.clone() 71 | x[0] = x[0] * (0.229 / 0.5) + (0.485 - 0.5) / 0.5 72 | x[1] = x[1] * (0.224 / 0.5) + (0.456 - 0.5) / 0.5 73 | x[2] = x[2] * (0.225 / 0.5) + (0.406 - 0.5) / 0.5 74 | # 299 x 299 x 3 75 | x = self.Conv2d_1a_3x3(x) 76 | # 149 x 149 x 32 77 | x = self.Conv2d_2a_3x3(x) 78 | # 147 x 147 x 32 79 | x = self.Conv2d_2b_3x3(x) 80 | # 147 x 147 x 64 81 | x = F.max_pool2d(x, kernel_size=3, stride=2) 82 | # 73 x 73 x 64 83 | x = self.Conv2d_3b_1x1(x) 84 | # 73 x 73 x 80 85 | x = self.Conv2d_4a_3x3(x) 86 | # 71 x 71 x 192 87 | x = F.max_pool2d(x, kernel_size=3, stride=2) 88 | # 35 x 35 x 192 89 | x = self.Mixed_5b(x) 90 | # 35 x 35 x 256 91 | x = self.Mixed_5c(x) 92 | # 35 x 35 x 288 93 | x = self.Mixed_5d(x) 94 | # 35 x 35 x 288 95 | x = self.Mixed_6a(x) 96 | # 17 x 17 x 768 97 | x = self.Mixed_6b(x) 98 | # 17 x 17 x 768 99 | x = self.Mixed_6c(x) 100 | # 17 x 17 x 768 101 | x = self.Mixed_6d(x) 102 | # 17 x 17 x 768 103 | x = self.Mixed_6e(x) 104 | # 17 x 17 x 768 105 | if self.training and self.aux_logits: 106 | aux = self.AuxLogits(x) 107 | # 17 x 17 x 768 108 | x = self.Mixed_7a(x) 109 | # 8 x 8 x 1280 110 | x = self.Mixed_7b(x) 111 | # 8 x 8 x 2048 112 | x = self.Mixed_7c(x) 113 | # 8 x 8 x 2048 114 | x = F.avg_pool2d(x, kernel_size=8) 115 | # 1 x 1 x 2048 116 | x = F.dropout(x, training=self.training) 117 | # 1 x 1 x 2048 118 | x = x.view(x.size(0), -1) 119 | # 2048 120 | x = self.group1(x) 121 | # 1000 (num_classes) 122 | if self.training and self.aux_logits: 123 | return x, aux 124 | return x 125 | 126 | 127 | class InceptionA(nn.Module): 128 | 129 | def __init__(self, in_channels, pool_features): 130 | super(InceptionA, self).__init__() 131 | self.branch1x1 = BasicConv2d(in_channels, 64, kernel_size=1) 132 | 133 | self.branch5x5_1 = BasicConv2d(in_channels, 48, kernel_size=1) 134 | self.branch5x5_2 = BasicConv2d(48, 64, kernel_size=5, padding=2) 135 | 136 | self.branch3x3dbl_1 = BasicConv2d(in_channels, 64, kernel_size=1) 137 | self.branch3x3dbl_2 = BasicConv2d(64, 96, kernel_size=3, padding=1) 138 | self.branch3x3dbl_3 = BasicConv2d(96, 96, kernel_size=3, padding=1) 139 | 140 | self.branch_pool = BasicConv2d(in_channels, pool_features, kernel_size=1) 141 | 142 | def forward(self, x): 143 | branch1x1 = self.branch1x1(x) 144 | 145 | branch5x5 = self.branch5x5_1(x) 146 | branch5x5 = self.branch5x5_2(branch5x5) 147 | 148 | branch3x3dbl = self.branch3x3dbl_1(x) 149 | branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) 150 | branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl) 151 | 152 | branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1) 153 | branch_pool = self.branch_pool(branch_pool) 154 | 155 | outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool] 156 | return torch.cat(outputs, 1) 157 | 158 | 159 | class InceptionB(nn.Module): 160 | 161 | def __init__(self, in_channels): 162 | super(InceptionB, self).__init__() 163 | self.branch3x3 = BasicConv2d(in_channels, 384, kernel_size=3, stride=2) 164 | 165 | self.branch3x3dbl_1 = BasicConv2d(in_channels, 64, kernel_size=1) 166 | self.branch3x3dbl_2 = BasicConv2d(64, 96, kernel_size=3, padding=1) 167 | self.branch3x3dbl_3 = BasicConv2d(96, 96, kernel_size=3, stride=2) 168 | 169 | def forward(self, x): 170 | branch3x3 = self.branch3x3(x) 171 | 172 | branch3x3dbl = self.branch3x3dbl_1(x) 173 | branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) 174 | branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl) 175 | 176 | branch_pool = F.max_pool2d(x, kernel_size=3, stride=2) 177 | 178 | outputs = [branch3x3, branch3x3dbl, branch_pool] 179 | return torch.cat(outputs, 1) 180 | 181 | 182 | class InceptionC(nn.Module): 183 | 184 | def __init__(self, in_channels, channels_7x7): 185 | super(InceptionC, self).__init__() 186 | self.branch1x1 = BasicConv2d(in_channels, 192, kernel_size=1) 187 | 188 | c7 = channels_7x7 189 | self.branch7x7_1 = BasicConv2d(in_channels, c7, kernel_size=1) 190 | self.branch7x7_2 = BasicConv2d(c7, c7, kernel_size=(1, 7), padding=(0, 3)) 191 | self.branch7x7_3 = BasicConv2d(c7, 192, kernel_size=(7, 1), padding=(3, 0)) 192 | 193 | self.branch7x7dbl_1 = BasicConv2d(in_channels, c7, kernel_size=1) 194 | self.branch7x7dbl_2 = BasicConv2d(c7, c7, kernel_size=(7, 1), padding=(3, 0)) 195 | self.branch7x7dbl_3 = BasicConv2d(c7, c7, kernel_size=(1, 7), padding=(0, 3)) 196 | self.branch7x7dbl_4 = BasicConv2d(c7, c7, kernel_size=(7, 1), padding=(3, 0)) 197 | self.branch7x7dbl_5 = BasicConv2d(c7, 192, kernel_size=(1, 7), padding=(0, 3)) 198 | 199 | self.branch_pool = BasicConv2d(in_channels, 192, kernel_size=1) 200 | 201 | def forward(self, x): 202 | branch1x1 = self.branch1x1(x) 203 | 204 | branch7x7 = self.branch7x7_1(x) 205 | branch7x7 = self.branch7x7_2(branch7x7) 206 | branch7x7 = self.branch7x7_3(branch7x7) 207 | 208 | branch7x7dbl = self.branch7x7dbl_1(x) 209 | branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl) 210 | branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl) 211 | branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl) 212 | branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl) 213 | 214 | branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1) 215 | branch_pool = self.branch_pool(branch_pool) 216 | 217 | outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool] 218 | return torch.cat(outputs, 1) 219 | 220 | 221 | class InceptionD(nn.Module): 222 | 223 | def __init__(self, in_channels): 224 | super(InceptionD, self).__init__() 225 | self.branch3x3_1 = BasicConv2d(in_channels, 192, kernel_size=1) 226 | self.branch3x3_2 = BasicConv2d(192, 320, kernel_size=3, stride=2) 227 | 228 | self.branch7x7x3_1 = BasicConv2d(in_channels, 192, kernel_size=1) 229 | self.branch7x7x3_2 = BasicConv2d(192, 192, kernel_size=(1, 7), padding=(0, 3)) 230 | self.branch7x7x3_3 = BasicConv2d(192, 192, kernel_size=(7, 1), padding=(3, 0)) 231 | self.branch7x7x3_4 = BasicConv2d(192, 192, kernel_size=3, stride=2) 232 | 233 | def forward(self, x): 234 | branch3x3 = self.branch3x3_1(x) 235 | branch3x3 = self.branch3x3_2(branch3x3) 236 | 237 | branch7x7x3 = self.branch7x7x3_1(x) 238 | branch7x7x3 = self.branch7x7x3_2(branch7x7x3) 239 | branch7x7x3 = self.branch7x7x3_3(branch7x7x3) 240 | branch7x7x3 = self.branch7x7x3_4(branch7x7x3) 241 | 242 | branch_pool = F.max_pool2d(x, kernel_size=3, stride=2) 243 | outputs = [branch3x3, branch7x7x3, branch_pool] 244 | return torch.cat(outputs, 1) 245 | 246 | 247 | class InceptionE(nn.Module): 248 | 249 | def __init__(self, in_channels): 250 | super(InceptionE, self).__init__() 251 | self.branch1x1 = BasicConv2d(in_channels, 320, kernel_size=1) 252 | 253 | self.branch3x3_1 = BasicConv2d(in_channels, 384, kernel_size=1) 254 | self.branch3x3_2a = BasicConv2d(384, 384, kernel_size=(1, 3), padding=(0, 1)) 255 | self.branch3x3_2b = BasicConv2d(384, 384, kernel_size=(3, 1), padding=(1, 0)) 256 | 257 | self.branch3x3dbl_1 = BasicConv2d(in_channels, 448, kernel_size=1) 258 | self.branch3x3dbl_2 = BasicConv2d(448, 384, kernel_size=3, padding=1) 259 | self.branch3x3dbl_3a = BasicConv2d(384, 384, kernel_size=(1, 3), padding=(0, 1)) 260 | self.branch3x3dbl_3b = BasicConv2d(384, 384, kernel_size=(3, 1), padding=(1, 0)) 261 | 262 | self.branch_pool = BasicConv2d(in_channels, 192, kernel_size=1) 263 | 264 | def forward(self, x): 265 | branch1x1 = self.branch1x1(x) 266 | 267 | branch3x3 = self.branch3x3_1(x) 268 | branch3x3 = [ 269 | self.branch3x3_2a(branch3x3), 270 | self.branch3x3_2b(branch3x3), 271 | ] 272 | branch3x3 = torch.cat(branch3x3, 1) 273 | 274 | branch3x3dbl = self.branch3x3dbl_1(x) 275 | branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) 276 | branch3x3dbl = [ 277 | self.branch3x3dbl_3a(branch3x3dbl), 278 | self.branch3x3dbl_3b(branch3x3dbl), 279 | ] 280 | branch3x3dbl = torch.cat(branch3x3dbl, 1) 281 | 282 | branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1) 283 | branch_pool = self.branch_pool(branch_pool) 284 | 285 | outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool] 286 | return torch.cat(outputs, 1) 287 | 288 | 289 | class InceptionAux(nn.Module): 290 | 291 | def __init__(self, in_channels, num_classes): 292 | super(InceptionAux, self).__init__() 293 | self.conv0 = BasicConv2d(in_channels, 128, kernel_size=1) 294 | self.conv1 = BasicConv2d(128, 768, kernel_size=5) 295 | self.conv1.stddev = 0.01 296 | 297 | fc = nn.Linear(768, num_classes) 298 | fc.stddev = 0.001 299 | 300 | self.group1 = nn.Sequential( 301 | OrderedDict([ 302 | ('fc', fc) 303 | ]) 304 | ) 305 | 306 | def forward(self, x): 307 | # 17 x 17 x 768 308 | x = F.avg_pool2d(x, kernel_size=5, stride=3) 309 | # 5 x 5 x 768 310 | x = self.conv0(x) 311 | # 5 x 5 x 128 312 | x = self.conv1(x) 313 | # 1 x 1 x 768 314 | x = x.view(x.size(0), -1) 315 | # 768 316 | x = self.group1(x) 317 | # 1000 318 | return x 319 | 320 | 321 | class BasicConv2d(nn.Module): 322 | 323 | def __init__(self, in_channels, out_channels, **kwargs): 324 | super(BasicConv2d, self).__init__() 325 | self.group1 = nn.Sequential( 326 | OrderedDict([ 327 | ('conv', nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)), 328 | ('bn', nn.BatchNorm2d(out_channels, eps=0.001)) 329 | ]) 330 | ) 331 | 332 | def forward(self, x): 333 | x = self.group1(x) 334 | return F.relu(x, inplace=True) 335 | -------------------------------------------------------------------------------- /imagenet/resnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import math 3 | from utee import misc 4 | from collections import OrderedDict 5 | 6 | __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 7 | 'resnet152'] 8 | 9 | 10 | model_urls = { 11 | 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 12 | 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 13 | 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 14 | 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 15 | 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 16 | } 17 | 18 | 19 | def conv3x3(in_planes, out_planes, stride=1): 20 | # "3x3 convolution with padding" 21 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) 22 | 23 | 24 | class BasicBlock(nn.Module): 25 | expansion = 1 26 | def __init__(self, inplanes, planes, stride=1, downsample=None): 27 | super(BasicBlock, self).__init__() 28 | m = OrderedDict() 29 | m['conv1'] = conv3x3(inplanes, planes, stride) 30 | m['bn1'] = nn.BatchNorm2d(planes) 31 | m['relu1'] = nn.ReLU(inplace=True) 32 | m['conv2'] = conv3x3(planes, planes) 33 | m['bn2'] = nn.BatchNorm2d(planes) 34 | self.group1 = nn.Sequential(m) 35 | 36 | self.relu= nn.Sequential(nn.ReLU(inplace=True)) 37 | self.downsample = downsample 38 | 39 | def forward(self, x): 40 | if self.downsample is not None: 41 | residual = self.downsample(x) 42 | else: 43 | residual = x 44 | 45 | out = self.group1(x) + residual 46 | 47 | out = self.relu(out) 48 | 49 | return out 50 | 51 | 52 | class Bottleneck(nn.Module): 53 | expansion = 4 54 | def __init__(self, inplanes, planes, stride=1, downsample=None): 55 | super(Bottleneck, self).__init__() 56 | m = OrderedDict() 57 | m['conv1'] = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 58 | m['bn1'] = nn.BatchNorm2d(planes) 59 | m['relu1'] = nn.ReLU(inplace=True) 60 | m['conv2'] = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) 61 | m['bn2'] = nn.BatchNorm2d(planes) 62 | m['relu2'] = nn.ReLU(inplace=True) 63 | m['conv3'] = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 64 | m['bn3'] = nn.BatchNorm2d(planes * 4) 65 | self.group1 = nn.Sequential(m) 66 | 67 | self.relu= nn.Sequential(nn.ReLU(inplace=True)) 68 | self.downsample = downsample 69 | 70 | def forward(self, x): 71 | if self.downsample is not None: 72 | residual = self.downsample(x) 73 | else: 74 | residual = x 75 | 76 | out = self.group1(x) + residual 77 | out = self.relu(out) 78 | 79 | return out 80 | 81 | 82 | class ResNet(nn.Module): 83 | def __init__(self, block, layers, num_classes=1000): 84 | self.inplanes = 64 85 | super(ResNet, self).__init__() 86 | 87 | m = OrderedDict() 88 | m['conv1'] = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) 89 | m['bn1'] = nn.BatchNorm2d(64) 90 | m['relu1'] = nn.ReLU(inplace=True) 91 | m['maxpool'] = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 92 | self.group1= nn.Sequential(m) 93 | 94 | self.layer1 = self._make_layer(block, 64, layers[0]) 95 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 96 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 97 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 98 | 99 | self.avgpool = nn.Sequential(nn.AvgPool2d(7)) 100 | 101 | self.group2 = nn.Sequential( 102 | OrderedDict([ 103 | ('fc', nn.Linear(512 * block.expansion, num_classes)) 104 | ]) 105 | ) 106 | 107 | for m in self.modules(): 108 | if isinstance(m, nn.Conv2d): 109 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 110 | m.weight.data.normal_(0, math.sqrt(2. / n)) 111 | elif isinstance(m, nn.BatchNorm2d): 112 | m.weight.data.fill_(1) 113 | m.bias.data.zero_() 114 | 115 | def _make_layer(self, block, planes, blocks, stride=1): 116 | downsample = None 117 | if stride != 1 or self.inplanes != planes * block.expansion: 118 | downsample = nn.Sequential( 119 | nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), 120 | nn.BatchNorm2d(planes * block.expansion), 121 | ) 122 | 123 | layers = [] 124 | layers.append(block(self.inplanes, planes, stride, downsample)) 125 | self.inplanes = planes * block.expansion 126 | for i in range(1, blocks): 127 | layers.append(block(self.inplanes, planes)) 128 | 129 | return nn.Sequential(*layers) 130 | 131 | def forward(self, x): 132 | x = self.group1(x) 133 | 134 | x = self.layer1(x) 135 | x = self.layer2(x) 136 | x = self.layer3(x) 137 | x = self.layer4(x) 138 | 139 | x = self.avgpool(x) 140 | x = x.view(x.size(0), -1) 141 | x = self.group2(x) 142 | 143 | return x 144 | 145 | 146 | def resnet18(pretrained=False, model_root=None, **kwargs): 147 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 148 | if pretrained: 149 | misc.load_state_dict(model, model_urls['resnet18'], model_root) 150 | return model 151 | 152 | 153 | def resnet34(pretrained=False, model_root=None, **kwargs): 154 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 155 | if pretrained: 156 | misc.load_state_dict(model, model_urls['resnet34'], model_root) 157 | return model 158 | 159 | 160 | def resnet50(pretrained=False, model_root=None, **kwargs): 161 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 162 | if pretrained: 163 | misc.load_state_dict(model, model_urls['resnet50'], model_root) 164 | return model 165 | 166 | 167 | def resnet101(pretrained=False, model_root=None, **kwargs): 168 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 169 | if pretrained: 170 | misc.load_state_dict(model, model_urls['resnet101'], model_root) 171 | return model 172 | 173 | 174 | def resnet152(pretrained=False, model_root=None, **kwargs): 175 | model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) 176 | if pretrained: 177 | misc.load_state_dict(model, model_urls['resnet152'], model_root) 178 | return model 179 | -------------------------------------------------------------------------------- /imagenet/squeezenet.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | import torch.nn as nn 4 | from utee import misc 5 | from collections import OrderedDict 6 | 7 | 8 | __all__ = ['SqueezeNet', 'squeezenet1_0', 'squeezenet1_1'] 9 | 10 | 11 | model_urls = { 12 | 'squeezenet1_0': 'https://download.pytorch.org/models/squeezenet1_0-a815701f.pth', 13 | 'squeezenet1_1': 'https://download.pytorch.org/models/squeezenet1_1-f364aa15.pth', 14 | } 15 | 16 | 17 | class Fire(nn.Module): 18 | 19 | def __init__(self, inplanes, squeeze_planes, 20 | expand1x1_planes, expand3x3_planes): 21 | super(Fire, self).__init__() 22 | self.inplanes = inplanes 23 | 24 | self.group1 = nn.Sequential( 25 | OrderedDict([ 26 | ('squeeze', nn.Conv2d(inplanes, squeeze_planes, kernel_size=1)), 27 | ('squeeze_activation', nn.ReLU(inplace=True)) 28 | ]) 29 | ) 30 | 31 | self.group2 = nn.Sequential( 32 | OrderedDict([ 33 | ('expand1x1', nn.Conv2d(squeeze_planes, expand1x1_planes, kernel_size=1)), 34 | ('expand1x1_activation', nn.ReLU(inplace=True)) 35 | ]) 36 | ) 37 | 38 | self.group3 = nn.Sequential( 39 | OrderedDict([ 40 | ('expand3x3', nn.Conv2d(squeeze_planes, expand3x3_planes, kernel_size=3, padding=1)), 41 | ('expand3x3_activation', nn.ReLU(inplace=True)) 42 | ]) 43 | ) 44 | 45 | def forward(self, x): 46 | x = self.group1(x) 47 | return torch.cat([self.group2(x),self.group3(x)], 1) 48 | 49 | 50 | class SqueezeNet(nn.Module): 51 | 52 | def __init__(self, version=1.0, num_classes=1000): 53 | super(SqueezeNet, self).__init__() 54 | if version not in [1.0, 1.1]: 55 | raise ValueError("Unsupported SqueezeNet version {version}:" 56 | "1.0 or 1.1 expected".format(version=version)) 57 | self.num_classes = num_classes 58 | if version == 1.0: 59 | self.features = nn.Sequential( 60 | nn.Conv2d(3, 96, kernel_size=7, stride=2), 61 | nn.ReLU(inplace=True), 62 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 63 | Fire(96, 16, 64, 64), 64 | Fire(128, 16, 64, 64), 65 | Fire(128, 32, 128, 128), 66 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 67 | Fire(256, 32, 128, 128), 68 | Fire(256, 48, 192, 192), 69 | Fire(384, 48, 192, 192), 70 | Fire(384, 64, 256, 256), 71 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 72 | Fire(512, 64, 256, 256), 73 | ) 74 | else: 75 | self.features = nn.Sequential( 76 | nn.Conv2d(3, 64, kernel_size=3, stride=2), 77 | nn.ReLU(inplace=True), 78 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 79 | Fire(64, 16, 64, 64), 80 | Fire(128, 16, 64, 64), 81 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 82 | Fire(128, 32, 128, 128), 83 | Fire(256, 32, 128, 128), 84 | nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), 85 | Fire(256, 48, 192, 192), 86 | Fire(384, 48, 192, 192), 87 | Fire(384, 64, 256, 256), 88 | Fire(512, 64, 256, 256), 89 | ) 90 | # Final convolution is initialized differently form the rest 91 | final_conv = nn.Conv2d(512, num_classes, kernel_size=1) 92 | self.classifier = nn.Sequential( 93 | nn.Dropout(p=0.5), 94 | final_conv, 95 | nn.ReLU(inplace=True), 96 | nn.AvgPool2d(13) 97 | ) 98 | 99 | for m in self.modules(): 100 | if isinstance(m, nn.Conv2d): 101 | gain = 2.0 102 | if m is final_conv: 103 | m.weight.data.normal_(0, 0.01) 104 | else: 105 | fan_in = m.kernel_size[0] * m.kernel_size[1] * m.in_channels 106 | u = math.sqrt(3.0 * gain / fan_in) 107 | m.weight.data.uniform_(-u, u) 108 | if m.bias is not None: 109 | m.bias.data.zero_() 110 | 111 | def forward(self, x): 112 | x = self.features(x) 113 | x = self.classifier(x) 114 | return x.view(x.size(0), self.num_classes) 115 | 116 | def squeezenet1_0(pretrained=False, model_root=None, **kwargs): 117 | r"""SqueezeNet model architecture from the `"SqueezeNet: AlexNet-level 118 | accuracy with 50x fewer parameters and <0.5MB model size" 119 | `_ paper. 120 | """ 121 | model = SqueezeNet(version=1.0, **kwargs) 122 | if pretrained: 123 | misc.load_state_dict(model, model_urls['squeezenet1_0'], model_root) 124 | return model 125 | 126 | 127 | def squeezenet1_1(pretrained=False, model_root=None, **kwargs): 128 | r"""SqueezeNet 1.1 model from the `official SqueezeNet repo 129 | `_. 130 | SqueezeNet 1.1 has 2.4x less computation and slightly fewer parameters 131 | than SqueezeNet 1.0, without sacrificing accuracy. 132 | """ 133 | model = SqueezeNet(version=1.1, **kwargs) 134 | if pretrained: 135 | misc.load_state_dict(model, model_urls['squeezenet1_1'], model_root) 136 | return model 137 | 138 | -------------------------------------------------------------------------------- /imagenet/vgg.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.utils.model_zoo as model_zoo 3 | import math 4 | 5 | 6 | __all__ = [ 7 | 'VGG', 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn', 8 | 'vgg19_bn', 'vgg19', 9 | ] 10 | 11 | 12 | model_urls = { 13 | 'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth', 14 | 'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth', 15 | 'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth', 16 | 'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth', 17 | } 18 | 19 | 20 | class VGG(nn.Module): 21 | 22 | def __init__(self, features, num_classes=1000): 23 | super(VGG, self).__init__() 24 | self.features = features 25 | self.classifier = nn.Sequential( 26 | nn.Linear(512 * 7 * 7, 4096), 27 | nn.ReLU(inplace=True), 28 | nn.Dropout(), 29 | nn.Linear(4096, 4096), 30 | nn.ReLU(inplace=True), 31 | nn.Dropout(), 32 | nn.Linear(4096, num_classes), 33 | ) 34 | self._initialize_weights() 35 | 36 | def forward(self, x): 37 | x = self.features(x) 38 | x = x.view(x.size(0), -1) 39 | x = self.classifier(x) 40 | return x 41 | 42 | def _initialize_weights(self): 43 | for m in self.modules(): 44 | if isinstance(m, nn.Conv2d): 45 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 46 | m.weight.data.normal_(0, math.sqrt(2. / n)) 47 | if m.bias is not None: 48 | m.bias.data.zero_() 49 | elif isinstance(m, nn.BatchNorm2d): 50 | m.weight.data.fill_(1) 51 | m.bias.data.zero_() 52 | elif isinstance(m, nn.Linear): 53 | n = m.weight.size(1) 54 | m.weight.data.normal_(0, 0.01) 55 | m.bias.data.zero_() 56 | 57 | 58 | def make_layers(cfg, batch_norm=False): 59 | layers = [] 60 | in_channels = 3 61 | for v in cfg: 62 | if v == 'M': 63 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 64 | else: 65 | conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) 66 | if batch_norm: 67 | layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] 68 | else: 69 | layers += [conv2d, nn.ReLU(inplace=True)] 70 | in_channels = v 71 | return nn.Sequential(*layers) 72 | 73 | 74 | cfg = { 75 | 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 76 | 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 77 | 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], 78 | 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'], 79 | } 80 | 81 | 82 | def vgg11(pretrained=False, model_root=None, **kwargs): 83 | """VGG 11-layer model (configuration "A")""" 84 | model = VGG(make_layers(cfg['A']), **kwargs) 85 | if pretrained: 86 | model.load_state_dict(model_zoo.load_url(model_urls['vgg11'], model_root)) 87 | return model 88 | 89 | 90 | def vgg11_bn(**kwargs): 91 | """VGG 11-layer model (configuration "A") with batch normalization""" 92 | kwargs.pop('model_root', None) 93 | return VGG(make_layers(cfg['A'], batch_norm=True), **kwargs) 94 | 95 | 96 | def vgg13(pretrained=False, model_root=None, **kwargs): 97 | """VGG 13-layer model (configuration "B")""" 98 | model = VGG(make_layers(cfg['B']), **kwargs) 99 | if pretrained: 100 | model.load_state_dict(model_zoo.load_url(model_urls['vgg13'], model_root)) 101 | return model 102 | 103 | 104 | def vgg13_bn(**kwargs): 105 | """VGG 13-layer model (configuration "B") with batch normalization""" 106 | kwargs.pop('model_root', None) 107 | return VGG(make_layers(cfg['B'], batch_norm=True), **kwargs) 108 | 109 | 110 | def vgg16(pretrained=False, model_root=None, **kwargs): 111 | """VGG 16-layer model (configuration "D")""" 112 | model = VGG(make_layers(cfg['D']), **kwargs) 113 | if pretrained: 114 | model.load_state_dict(model_zoo.load_url(model_urls['vgg16'], model_root)) 115 | return model 116 | 117 | 118 | def vgg16_bn(**kwargs): 119 | """VGG 16-layer model (configuration "D") with batch normalization""" 120 | kwargs.pop('model_root', None) 121 | return VGG(make_layers(cfg['D'], batch_norm=True), **kwargs) 122 | 123 | 124 | def vgg19(pretrained=False, model_root=None, **kwargs): 125 | """VGG 19-layer model (configuration "E")""" 126 | model = VGG(make_layers(cfg['E']), **kwargs) 127 | if pretrained: 128 | model.load_state_dict(model_zoo.load_url(model_urls['vgg19'], model_root)) 129 | return model 130 | 131 | 132 | def vgg19_bn(**kwargs): 133 | """VGG 19-layer model (configuration 'E') with batch normalization""" 134 | kwargs.pop('model_root', None) 135 | return VGG(make_layers(cfg['E'], batch_norm=True), **kwargs) 136 | -------------------------------------------------------------------------------- /mnist/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aaron-xichen/pytorch-playground/5d476295f1c1564c8074eb885bb6138d6bd09889/mnist/__init__.py -------------------------------------------------------------------------------- /mnist/dataset.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import DataLoader 2 | import torch 3 | from torchvision import datasets, transforms 4 | import os 5 | 6 | def get(batch_size, data_root='/tmp/public_dataset/pytorch', train=True, val=True, **kwargs): 7 | data_root = os.path.expanduser(os.path.join(data_root, 'mnist-data')) 8 | kwargs.pop('input_size', None) 9 | num_workers = kwargs.setdefault('num_workers', 1) 10 | print("Building MNIST data loader with {} workers".format(num_workers)) 11 | ds = [] 12 | if train: 13 | train_loader = torch.utils.data.DataLoader( 14 | datasets.MNIST(root=data_root, train=True, download=True, 15 | transform=transforms.Compose([ 16 | transforms.ToTensor(), 17 | transforms.Normalize((0.1307,), (0.3081,)) 18 | ])), 19 | batch_size=batch_size, shuffle=True, **kwargs) 20 | ds.append(train_loader) 21 | if val: 22 | test_loader = torch.utils.data.DataLoader( 23 | datasets.MNIST(root=data_root, train=False, download=True, 24 | transform=transforms.Compose([ 25 | transforms.ToTensor(), 26 | transforms.Normalize((0.1307,), (0.3081,)) 27 | ])), 28 | batch_size=batch_size, shuffle=True, **kwargs) 29 | ds.append(test_loader) 30 | ds = ds[0] if len(ds) == 1 else ds 31 | return ds 32 | 33 | -------------------------------------------------------------------------------- /mnist/model.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from collections import OrderedDict 3 | import torch.utils.model_zoo as model_zoo 4 | from utee import misc 5 | print = misc.logger.info 6 | 7 | model_urls = { 8 | 'mnist': 'http://ml.cs.tsinghua.edu.cn/~chenxi/pytorch-models/mnist-b07bb66b.pth' 9 | } 10 | 11 | class MLP(nn.Module): 12 | def __init__(self, input_dims, n_hiddens, n_class): 13 | super(MLP, self).__init__() 14 | assert isinstance(input_dims, int), 'Please provide int for input_dims' 15 | self.input_dims = input_dims 16 | current_dims = input_dims 17 | layers = OrderedDict() 18 | 19 | if isinstance(n_hiddens, int): 20 | n_hiddens = [n_hiddens] 21 | else: 22 | n_hiddens = list(n_hiddens) 23 | for i, n_hidden in enumerate(n_hiddens): 24 | layers['fc{}'.format(i+1)] = nn.Linear(current_dims, n_hidden) 25 | layers['relu{}'.format(i+1)] = nn.ReLU() 26 | layers['drop{}'.format(i+1)] = nn.Dropout(0.2) 27 | current_dims = n_hidden 28 | layers['out'] = nn.Linear(current_dims, n_class) 29 | 30 | self.model= nn.Sequential(layers) 31 | print(self.model) 32 | 33 | def forward(self, input): 34 | input = input.view(input.size(0), -1) 35 | assert input.size(1) == self.input_dims 36 | return self.model.forward(input) 37 | 38 | def mnist(input_dims=784, n_hiddens=[256, 256], n_class=10, pretrained=None): 39 | model = MLP(input_dims, n_hiddens, n_class) 40 | if pretrained is not None: 41 | m = model_zoo.load_url(model_urls['mnist']) 42 | state_dict = m.state_dict() if isinstance(m, nn.Module) else m 43 | assert isinstance(state_dict, (dict, OrderedDict)), type(state_dict) 44 | model.load_state_dict(state_dict) 45 | return model 46 | 47 | -------------------------------------------------------------------------------- /mnist/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import time 4 | 5 | from utee import misc 6 | import torch 7 | import torch.nn.functional as F 8 | import torch.optim as optim 9 | from torch.autograd import Variable 10 | 11 | import dataset 12 | import model 13 | 14 | parser = argparse.ArgumentParser(description='PyTorch MNIST Example') 15 | parser.add_argument('--wd', type=float, default=0.0001, help='weight decay') 16 | parser.add_argument('--batch_size', type=int, default=200, help='input batch size for training (default: 64)') 17 | parser.add_argument('--epochs', type=int, default=40, help='number of epochs to train (default: 10)') 18 | parser.add_argument('--lr', type=float, default=0.01, help='learning rate (default: 1e-3)') 19 | parser.add_argument('--gpu', default=None, help='index of gpus to use') 20 | parser.add_argument('--ngpu', type=int, default=1, help='number of gpus to use') 21 | parser.add_argument('--seed', type=int, default=117, help='random seed (default: 1)') 22 | parser.add_argument('--log_interval', type=int, default=100, help='how many batches to wait before logging training status') 23 | parser.add_argument('--test_interval', type=int, default=5, help='how many epochs to wait before another test') 24 | parser.add_argument('--logdir', default='log/default', help='folder to save to the log') 25 | parser.add_argument('--data_root', default='/tmp/public_dataset/pytorch/', help='folder to save the model') 26 | parser.add_argument('--decreasing_lr', default='80,120', help='decreasing strategy') 27 | args = parser.parse_args() 28 | args.logdir = os.path.join(os.path.dirname(__file__), args.logdir) 29 | misc.logger.init(args.logdir, 'train_log') 30 | print = misc.logger.info 31 | 32 | # select gpu 33 | args.gpu = misc.auto_select_gpu(utility_bound=0, num_gpu=args.ngpu, selected_gpus=args.gpu) 34 | args.ngpu = len(args.gpu) 35 | 36 | # logger 37 | misc.ensure_dir(args.logdir) 38 | print("=================FLAGS==================") 39 | for k, v in args.__dict__.items(): 40 | print('{}: {}'.format(k, v)) 41 | print("========================================") 42 | 43 | # seed 44 | args.cuda = torch.cuda.is_available() 45 | torch.manual_seed(args.seed) 46 | if args.cuda: 47 | torch.cuda.manual_seed(args.seed) 48 | 49 | # data loader 50 | train_loader, test_loader = dataset.get(batch_size=args.batch_size, data_root=args.data_root, num_workers=1) 51 | 52 | # model 53 | model = model.mnist(input_dims=784, n_hiddens=[256, 256], n_class=10) 54 | model = torch.nn.DataParallel(model, device_ids= range(args.ngpu)) 55 | if args.cuda: 56 | model.cuda() 57 | 58 | # optimizer 59 | optimizer = optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.wd, momentum=0.9) 60 | decreasing_lr = list(map(int, args.decreasing_lr.split(','))) 61 | print('decreasing_lr: ' + str(decreasing_lr)) 62 | best_acc, old_file = 0, None 63 | t_begin = time.time() 64 | try: 65 | # ready to go 66 | for epoch in range(args.epochs): 67 | model.train() 68 | if epoch in decreasing_lr: 69 | optimizer.param_groups[0]['lr'] *= 0.1 70 | for batch_idx, (data, target) in enumerate(train_loader): 71 | indx_target = target.clone() 72 | if args.cuda: 73 | data, target = data.cuda(), target.cuda() 74 | data, target = Variable(data), Variable(target) 75 | 76 | optimizer.zero_grad() 77 | output = model(data) 78 | loss = F.cross_entropy(output, target) 79 | loss.backward() 80 | optimizer.step() 81 | 82 | if batch_idx % args.log_interval == 0 and batch_idx > 0: 83 | pred = output.data.max(1)[1] # get the index of the max log-probability 84 | correct = pred.cpu().eq(indx_target).sum() 85 | acc = correct * 1.0 / len(data) 86 | print('Train Epoch: {} [{}/{}] Loss: {:.6f} Acc: {:.4f} lr: {:.2e}'.format( 87 | epoch, batch_idx * len(data), len(train_loader.dataset), 88 | loss.data, acc, optimizer.param_groups[0]['lr'])) 89 | 90 | elapse_time = time.time() - t_begin 91 | speed_epoch = elapse_time / (epoch + 1) 92 | speed_batch = speed_epoch / len(train_loader) 93 | eta = speed_epoch * args.epochs - elapse_time 94 | print("Elapsed {:.2f}s, {:.2f} s/epoch, {:.2f} s/batch, ets {:.2f}s".format( 95 | elapse_time, speed_epoch, speed_batch, eta)) 96 | misc.model_snapshot(model, os.path.join(args.logdir, 'latest.pth')) 97 | 98 | if epoch % args.test_interval == 0: 99 | model.eval() 100 | test_loss = 0 101 | correct = 0 102 | for data, target in test_loader: 103 | indx_target = target.clone() 104 | if args.cuda: 105 | data, target = data.cuda(), target.cuda() 106 | data, target = Variable(data, volatile=True), Variable(target) 107 | output = model(data) 108 | test_loss += F.cross_entropy(output, target).data 109 | pred = output.data.max(1)[1] # get the index of the max log-probability 110 | correct += pred.cpu().eq(indx_target).sum() 111 | 112 | test_loss = test_loss / len(test_loader) # average over number of mini-batch 113 | acc = 100. * correct / len(test_loader.dataset) 114 | print('\tTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format( 115 | test_loss, correct, len(test_loader.dataset), acc)) 116 | if acc > best_acc: 117 | new_file = os.path.join(args.logdir, 'best-{}.pth'.format(epoch)) 118 | misc.model_snapshot(model, new_file, old_file=old_file, verbose=True) 119 | best_acc = acc 120 | old_file = new_file 121 | except Exception as e: 122 | import traceback 123 | traceback.print_exc() 124 | finally: 125 | print("Total Elapse: {:.2f}, Best Result: {:.3f}%".format(time.time()-t_begin, best_acc)) 126 | 127 | 128 | -------------------------------------------------------------------------------- /quantize.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from utee import misc, quant, selector 3 | import torch 4 | import torch.backends.cudnn as cudnn 5 | cudnn.benchmark =True 6 | from collections import OrderedDict 7 | 8 | def main(): 9 | parser = argparse.ArgumentParser(description='PyTorch SVHN Example') 10 | parser.add_argument('--type', default='cifar10', help='|'.join(selector.known_models)) 11 | parser.add_argument('--quant_method', default='linear', help='linear|minmax|log|tanh') 12 | parser.add_argument('--batch_size', type=int, default=100, help='input batch size for training (default: 64)') 13 | parser.add_argument('--gpu', default=None, help='index of gpus to use') 14 | parser.add_argument('--ngpu', type=int, default=8, help='number of gpus to use') 15 | parser.add_argument('--seed', type=int, default=117, help='random seed (default: 1)') 16 | parser.add_argument('--model_root', default='~/.torch/models/', help='folder to save the model') 17 | parser.add_argument('--data_root', default='/data/public_dataset/pytorch/', help='folder to save the model') 18 | parser.add_argument('--logdir', default='log/default', help='folder to save to the log') 19 | 20 | parser.add_argument('--input_size', type=int, default=224, help='input size of image') 21 | parser.add_argument('--n_sample', type=int, default=20, help='number of samples to infer the scaling factor') 22 | parser.add_argument('--param_bits', type=int, default=8, help='bit-width for parameters') 23 | parser.add_argument('--bn_bits', type=int, default=32, help='bit-width for running mean and std') 24 | parser.add_argument('--fwd_bits', type=int, default=8, help='bit-width for layer output') 25 | parser.add_argument('--overflow_rate', type=float, default=0.0, help='overflow rate') 26 | args = parser.parse_args() 27 | 28 | args.gpu = misc.auto_select_gpu(utility_bound=0, num_gpu=args.ngpu, selected_gpus=args.gpu) 29 | args.ngpu = len(args.gpu) 30 | misc.ensure_dir(args.logdir) 31 | args.model_root = misc.expand_user(args.model_root) 32 | args.data_root = misc.expand_user(args.data_root) 33 | args.input_size = 299 if 'inception' in args.type else args.input_size 34 | assert args.quant_method in ['linear', 'minmax', 'log', 'tanh'] 35 | print("=================FLAGS==================") 36 | for k, v in args.__dict__.items(): 37 | print('{}: {}'.format(k, v)) 38 | print("========================================") 39 | 40 | assert torch.cuda.is_available(), 'no cuda' 41 | torch.manual_seed(args.seed) 42 | torch.cuda.manual_seed(args.seed) 43 | 44 | # load model and dataset fetcher 45 | model_raw, ds_fetcher, is_imagenet = selector.select(args.type, model_root=args.model_root) 46 | args.ngpu = args.ngpu if is_imagenet else 1 47 | 48 | # quantize parameters 49 | if args.param_bits < 32: 50 | state_dict = model_raw.state_dict() 51 | state_dict_quant = OrderedDict() 52 | sf_dict = OrderedDict() 53 | for k, v in state_dict.items(): 54 | if 'running' in k: 55 | if args.bn_bits >=32: 56 | print("Ignoring {}".format(k)) 57 | state_dict_quant[k] = v 58 | continue 59 | else: 60 | bits = args.bn_bits 61 | else: 62 | bits = args.param_bits 63 | 64 | if args.quant_method == 'linear': 65 | sf = bits - 1. - quant.compute_integral_part(v, overflow_rate=args.overflow_rate) 66 | v_quant = quant.linear_quantize(v, sf, bits=bits) 67 | elif args.quant_method == 'log': 68 | v_quant = quant.log_minmax_quantize(v, bits=bits) 69 | elif args.quant_method == 'minmax': 70 | v_quant = quant.min_max_quantize(v, bits=bits) 71 | else: 72 | v_quant = quant.tanh_quantize(v, bits=bits) 73 | state_dict_quant[k] = v_quant 74 | print(k, bits) 75 | model_raw.load_state_dict(state_dict_quant) 76 | 77 | # quantize forward activation 78 | if args.fwd_bits < 32: 79 | model_raw = quant.duplicate_model_with_quant(model_raw, bits=args.fwd_bits, overflow_rate=args.overflow_rate, 80 | counter=args.n_sample, type=args.quant_method) 81 | print(model_raw) 82 | val_ds_tmp = ds_fetcher(10, data_root=args.data_root, train=False, input_size=args.input_size) 83 | misc.eval_model(model_raw, val_ds_tmp, ngpu=1, n_sample=args.n_sample, is_imagenet=is_imagenet) 84 | 85 | # eval model 86 | val_ds = ds_fetcher(args.batch_size, data_root=args.data_root, train=False, input_size=args.input_size) 87 | acc1, acc5 = misc.eval_model(model_raw, val_ds, ngpu=args.ngpu, is_imagenet=is_imagenet) 88 | 89 | # print sf 90 | print(model_raw) 91 | res_str = "type={}, quant_method={}, param_bits={}, bn_bits={}, fwd_bits={}, overflow_rate={}, acc1={:.4f}, acc5={:.4f}".format( 92 | args.type, args.quant_method, args.param_bits, args.bn_bits, args.fwd_bits, args.overflow_rate, acc1, acc5) 93 | print(res_str) 94 | with open('acc1_acc5.txt', 'a') as f: 95 | f.write(res_str + '\n') 96 | 97 | 98 | if __name__ == '__main__': 99 | main() 100 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | Pillow==6.1 2 | torchvision==0.4.2 3 | tqdm==4.41.1 4 | opencv-python==4.1.2.30 5 | joblib==0.14.1 6 | -------------------------------------------------------------------------------- /roadmap_zh.md: -------------------------------------------------------------------------------- 1 | # 定点化Roadmap 2 | 首先定点化的setting分好几种，主要如下所示 (w代表weight，a代表activation，g代表gradient) 3 | 4 | 最近两年的目前有13篇直接相关的论文，截止到2016年7月 5 | 6 | ## float转化为定点版本，不允许fine-tune 7 | - w定点，a浮点 8 | - Resiliency of Deep Neural Networks under Quantization [Wongyong Sung, Sungho Shin, 2016.01.07, ICLR2016] {5bit在CIFAR10上恢复正确率} 9 | - Fixed Point Quantization of Deep Convolutional Networks [Darryl D.Lin, Sachin S.Talathi, 2016.06.02] {每层定点化策略不同，解析解求出} 10 | - w+a定点 11 | - Hardware-oriented approximation of convolutional neural networks [Philipp Gysel, Mohammad Motamedi, ICLR 2016 Workshop] {ImageNet上8bit-8bit掉0.9%，AlexNet} 12 | - Energy-Efficient ConvNets Through Approximate Computing [Bert Moons, KU leuven, 2016.03.22] {结合硬件的trick可以在ImageNet上4-10bit} 13 | - Going Deeper with Embedded FPGA Platform for Convolutional Neural Network [Jiantao Qiu, Jie Wang, FPGA2016]{ImageNet上8bit-8bit掉1%，AlexNet} 14 | 15 | ## float转化为定点版本，允许fine-tune 16 | - fine-tune整个网络 17 | - w定点，a+g浮点 18 | - Resiliency of Deep Neural Networks under Quantization [Wongyong Sung, Sungho Shin, 2016.01.07, ICLR2016] {2bit即三值网络在CIFAR10上恢复正确率} 19 | - w+a定点，g浮点 20 | - Fixed Point Quantization of Deep Convolutional Networks [Darryl D.Lin, Sachin S.Talathi, 2016.06.02] {每层定点化策略不同，解析解求出，CIFAR10上fine-tune后4bit-4bit掉1.32%} 21 | - w+a+g定点 22 | - Overcoming Challenges in Fixed Point Training of Deep Convolutional Networks [Darryl D.Lin, Sachin S. Talathi, Qualcomm Research，2016.07.08] {无随机rounding，ImageNet上4bit-16bit-16bit掉7.2%，a和g再小就不收敛} 23 | - DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients [Shuchang Zhou, Zekun Ni, 2016.06.20] {1bit-2bit-4bit, 第一层和最后一层没有量化，ImageNet上掉5.2%} 24 | - fine-tune最高几层 25 | - w+a+g定点 26 | - Overcoming Challenges in Fixed Point Training of Deep Convolutional Networks [Darryl D.Lin, Sachin S. Talathi, Qualcomm Research，2016.07.08] {无随机rounding，ImageNet上4bit-4bit-4bit掉23.3%} 27 | - 分阶段地从低层到高层fine-tune网络 28 | - w+a+g定点 29 | - Overcoming Challenges in Fixed Point Training of Deep Convolutional Networks [Darryl D.Lin, Sachin S. Talathi, Qualcomm Research，2016.07.08] {无随机rounding，ImageNet上4bit-4bit-4bit Top5掉11.5%} 30 | 31 | ## 直接定点从头开始训练 32 | - w定点，a+g浮点 33 | - 二值网络 34 | - BinaryConnect: Training Deep Neural Networks with binary weights during propagations [Matthieu Courbariaux, Yoshua Bengio, 2015.11.02, NIPS] {CIFAR10上8.27%, state-of-art} 35 | - XNOR-Net: ImageNet Classification Using Binary Convolutional Neural Networks [Mohammad Rastegari, Washington University, 2016.03.16] {ImageNet上39.2%，掉2.8%, AlexNet} 36 | - 三值网络 37 | - Ternary Weight Networks [Fengfu Li, Bin Liu, UCAS, China, 2016.05.16] {ImageNet掉2.3%, ResNet-18B} 38 | - Trained Ternary Quantization [Chenzhuo Zhu, Song Han, Huizi Mao, William J. Dally, ICLR2017] {ResNet上效果更佳} 39 | - w+a定点，g浮点 40 | - 二值网络 41 | - Binarized Neural Networks: Training Neural Networks with Weights and Activations Constrained to +1 or −1 [Matthieu Courbariaux, Yoshua Bengio, 2016.03.17] {CIFAR10上10.15%} 42 | - XNOR-Net: ImageNet Classification Using Binary Convolutional Neural Networks [Mohammad Rastegari, Washington University, 2016.03.16] {ImageNet上55.8%， 掉12.4%} 43 | - w+a+g定点 44 | - Deep Learning with Limited Numerical Precision [ Suyog Gupta, Ankur Agrawal, IBM, 2015.02.09] {随机rounding技巧，CIFAR10上16bit+16bit+16bit复现正确率} 45 | - Training deep neural networks with low precision multiplications [Matthieu Courbariaux, Yoshua Bengio, ICLR 2015 Workshop] {CIFAR10上10bit+10bit+12bit复现正确率} 46 | - DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients [Shuchang Zhou, Zekun Ni, 2016.06.20] {1bit-2bit-4bit, 第一层和最后一层没有量化，ImageNet上掉8.8%} 47 | - Quantized Neural Networks: Training Neural Networks with Low Precision Weights and Activations [Itay Hubara, Matthieu Courbariaux, 2016.09.22]{1bit-2bit-6bit，ImageNet上超过DoReFa 0.33%} -------------------------------------------------------------------------------- /script/convert.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import tqdm 4 | from utee import misc 5 | import argparse 6 | import cv2 7 | import joblib 8 | 9 | parser = argparse.ArgumentParser(description='Extract the ILSVRC2012 val dataset') 10 | parser.add_argument('--in_file', default='val224_compressed.pkl', help='input file path') 11 | parser.add_argument('--out_root', default='/data/public_dataset/pytorch/imagenet-data/', help='output file path') 12 | args = parser.parse_args() 13 | 14 | d = misc.load_pickle(args.in_file) 15 | assert len(d['data']) == 50000, len(d['data']) 16 | assert len(d['target']) == 50000, len(d['target']) 17 | 18 | 19 | data299 = [] 20 | for img, target in tqdm.tqdm(zip(d['data'], d['target']), total=50000): 21 | img224 = misc.str2img(img) 22 | img299 = cv2.resize(img224, (299, 299)) 23 | data299.append(img299) 24 | 25 | data_dict299 = dict( 26 | data = np.array(data299).transpose(0, 3, 1, 2), 27 | target = d['target'] 28 | ) 29 | if not os.path.exists(args.out_root): 30 | os.makedirs(args.out_root) 31 | joblib.dump(data_dict299, os.path.join(args.out_root, 'val299.pkl')) 32 | 33 | data299.clear() 34 | data_dict299.clear() 35 | 36 | data224 = [] 37 | for img, target in tqdm.tqdm(zip(d['data'], d['target']), total=50000): 38 | img224 = misc.str2img(img) 39 | data224.append(img224) 40 | 41 | data_dict224 = dict( 42 | data = np.array(data224).transpose(0, 3, 1, 2), 43 | target = d['target'] 44 | ) 45 | joblib.dump(data_dict224, os.path.join(args.out_root, 'val224.pkl')) 46 | 47 | 48 | 49 | 50 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | with open("requirements.txt") as requirements_file: 4 | REQUIREMENTS = requirements_file.readlines() 5 | 6 | setup( 7 | name="pytorch-playground", 8 | version="1.0.0", 9 | author='Aaron Chen', 10 | author_email='aaron.xichen@gmail.com', 11 | packages=find_packages(), 12 | entry_points = { 13 | 'console_scripts': [ 14 | 'quantize=quantize:main', 15 | ] 16 | }, 17 | install_requires=REQUIREMENTS, 18 | 19 | ) 20 | -------------------------------------------------------------------------------- /stl10/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aaron-xichen/pytorch-playground/5d476295f1c1564c8074eb885bb6138d6bd09889/stl10/__init__.py -------------------------------------------------------------------------------- /stl10/dataset.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torchvision import datasets, transforms 3 | from torch.utils.data import DataLoader 4 | from IPython import embed 5 | import os 6 | 7 | def get(batch_size, data_root='/mnt/local0/public_dataset/pytorch/', train=True, val=True, **kwargs): 8 | data_root = os.path.expanduser(os.path.join(data_root, 'stl10-data')) 9 | num_workers = kwargs.setdefault('num_workers', 1) 10 | kwargs.pop('input_size', None) 11 | print("Building STL10 data loader with {} workers".format(num_workers)) 12 | ds = [] 13 | if train: 14 | train_loader = torch.utils.data.DataLoader( 15 | datasets.STL10( 16 | root=data_root, split='train', download=True, 17 | transform=transforms.Compose([ 18 | transforms.Pad(4), 19 | transforms.RandomCrop(96), 20 | transforms.RandomHorizontalFlip(), 21 | transforms.ToTensor(), 22 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 23 | ])), 24 | batch_size=batch_size, shuffle=True, **kwargs) 25 | ds.append(train_loader) 26 | 27 | if val: 28 | test_loader = torch.utils.data.DataLoader( 29 | datasets.STL10( 30 | root=data_root, split='test', download=True, 31 | transform=transforms.Compose([ 32 | transforms.ToTensor(), 33 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 34 | ])), 35 | batch_size=batch_size, shuffle=False, **kwargs) 36 | ds.append(test_loader) 37 | 38 | ds = ds[0] if len(ds) == 1 else ds 39 | return ds 40 | 41 | if __name__ == '__main__': 42 | train_ds, test_ds = get(200, num_workers=1) 43 | for data, target in train_ds: 44 | print("~~") 45 | -------------------------------------------------------------------------------- /stl10/model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.utils.model_zoo as model_zoo 4 | import os 5 | from utee import misc 6 | from collections import OrderedDict 7 | print = misc.logger.info 8 | 9 | model_urls = { 10 | 'stl10': 'http://ml.cs.tsinghua.edu.cn/~chenxi/pytorch-models/stl10-866321e9.pth', 11 | } 12 | 13 | class SVHN(nn.Module): 14 | def __init__(self, features, n_channel, num_classes): 15 | super(SVHN, self).__init__() 16 | assert isinstance(features, nn.Sequential), type(features) 17 | self.features = features 18 | self.classifier = nn.Sequential( 19 | nn.Linear(n_channel, num_classes) 20 | ) 21 | print(self.features) 22 | print(self.classifier) 23 | 24 | def forward(self, x): 25 | x = self.features(x) 26 | x = x.view(x.size(0), -1) 27 | x = self.classifier(x) 28 | return x 29 | 30 | def make_layers(cfg, batch_norm=False): 31 | layers = [] 32 | in_channels = 3 33 | for i, v in enumerate(cfg): 34 | if v == 'M': 35 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 36 | else: 37 | padding = v[1] if isinstance(v, tuple) else 1 38 | out_channels = v[0] if isinstance(v, tuple) else v 39 | conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=padding) 40 | if batch_norm: 41 | layers += [conv2d, nn.BatchNorm2d(out_channels, affine=False), nn.ReLU()] 42 | else: 43 | layers += [conv2d, nn.ReLU()] 44 | in_channels = out_channels 45 | return nn.Sequential(*layers) 46 | 47 | def stl10(n_channel, pretrained=None): 48 | cfg = [ 49 | n_channel, 'M', 50 | 2*n_channel, 'M', 51 | 4*n_channel, 'M', 52 | 4*n_channel, 'M', 53 | (8*n_channel, 0), (8*n_channel, 0), 'M' 54 | ] 55 | layers = make_layers(cfg, batch_norm=True) 56 | model = SVHN(layers, n_channel=8*n_channel, num_classes=10) 57 | if pretrained is not None: 58 | m = model_zoo.load_url(model_urls['stl10']) 59 | state_dict = m.state_dict() if isinstance(m, nn.Module) else m 60 | assert isinstance(state_dict, (dict, OrderedDict)), type(state_dict) 61 | model.load_state_dict(state_dict) 62 | return model 63 | 64 | 65 | -------------------------------------------------------------------------------- /stl10/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import time 4 | from utee import misc 5 | 6 | import torch 7 | import torch.nn.functional as F 8 | import torch.optim as optim 9 | from torch.autograd import Variable 10 | 11 | import dataset 12 | import model 13 | from IPython import embed 14 | 15 | parser = argparse.ArgumentParser(description='PyTorch SVHN Example') 16 | parser.add_argument('--channel', type=int, default=32, help='first conv channel (default: 32)') 17 | parser.add_argument('--wd', type=float, default=0.00, help='weight decay') 18 | parser.add_argument('--batch_size', type=int, default=200, help='input batch size for training (default: 64)') 19 | parser.add_argument('--epochs', type=int, default=150, help='number of epochs to train (default: 10)') 20 | parser.add_argument('--lr', type=float, default=0.001, help='learning rate (default: 1e-3)') 21 | parser.add_argument('--gpu', default=None, help='index of gpus to use') 22 | parser.add_argument('--ngpu', type=int, default=2, help='number of gpus to use') 23 | parser.add_argument('--seed', type=int, default=117, help='random seed (default: 1)') 24 | parser.add_argument('--log_interval', type=int, default=20, help='how many batches to wait before logging training status') 25 | parser.add_argument('--test_interval', type=int, default=5, help='how many epochs to wait before another test') 26 | parser.add_argument('--logdir', default='log/default', help='folder to save to the log') 27 | parser.add_argument('--decreasing_lr', default='80,120', help='decreasing strategy') 28 | args = parser.parse_args() 29 | args.logdir = os.path.join(os.path.dirname(__file__), args.logdir) 30 | misc.logger.init(args.logdir, 'train_log') 31 | print = misc.logger.info 32 | 33 | # select gpu 34 | args.gpu = misc.auto_select_gpu(utility_bound=0, num_gpu=args.ngpu, selected_gpus=args.gpu) 35 | args.ngpu = len(args.gpu) 36 | 37 | # logger 38 | misc.ensure_dir(args.logdir) 39 | print("=================FLAGS==================") 40 | for k, v in args.__dict__.items(): 41 | print('{}: {}'.format(k, v)) 42 | print("========================================") 43 | 44 | # seed 45 | args.cuda = torch.cuda.is_available() 46 | torch.manual_seed(args.seed) 47 | if args.cuda: 48 | torch.cuda.manual_seed(args.seed) 49 | 50 | # data loader and model 51 | train_loader, test_loader = dataset.get(batch_size=args.batch_size, num_workers=1) 52 | model = model.stl10(n_channel=args.channel) 53 | model = torch.nn.DataParallel(model, device_ids= range(args.ngpu)) 54 | if args.cuda: 55 | model.cuda() 56 | 57 | # optimizer 58 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd) 59 | decreasing_lr = list(map(int, args.decreasing_lr.split(','))) 60 | print('decreasing_lr: ' + str(decreasing_lr)) 61 | best_acc, old_file = 0, None 62 | t_begin = time.time() 63 | try: 64 | # ready to go 65 | for epoch in range(args.epochs): 66 | model.train() 67 | if epoch in decreasing_lr: 68 | optimizer.param_groups[0]['lr'] *= 0.1 69 | for batch_idx, (data, target) in enumerate(train_loader): 70 | indx_target = target.clone() 71 | if args.cuda: 72 | data, target = data.cuda(), target.cuda() 73 | data, target = Variable(data), Variable(target) 74 | 75 | optimizer.zero_grad() 76 | output = model(data) 77 | loss = F.cross_entropy(output, target) 78 | loss.backward() 79 | optimizer.step() 80 | 81 | if batch_idx % args.log_interval == 0 and batch_idx > 0: 82 | pred = output.data.max(1)[1] # get the index of the max log-probability 83 | correct = pred.cpu().eq(indx_target).sum() 84 | acc = correct * 1.0 / len(data) 85 | print('Train Epoch: {} [{}/{}] Loss: {:.6f} Acc: {:.4f} lr: {:.2e}'.format( 86 | epoch, batch_idx * len(data), len(train_loader.dataset), 87 | loss.data[0], acc, optimizer.param_groups[0]['lr'])) 88 | 89 | elapse_time = time.time() - t_begin 90 | speed_epoch = elapse_time / (epoch + 1) 91 | speed_batch = speed_epoch / len(train_loader) 92 | eta = speed_epoch * args.epochs - elapse_time 93 | print("Elapsed {:.2f}s, {:.2f} s/epoch, {:.2f} s/batch, ets {:.2f}s".format( 94 | elapse_time, speed_epoch, speed_batch, eta)) 95 | misc.model_snapshot(model, os.path.join(args.logdir, 'latest.pth')) 96 | 97 | if epoch % args.test_interval == 0: 98 | model.eval() 99 | test_loss = 0 100 | correct = 0 101 | for data, target in test_loader: 102 | indx_target = target.clone() 103 | if args.cuda: 104 | data, target = data.cuda(), target.cuda().long().squeeze() 105 | data, target = Variable(data, volatile=True), Variable(target) 106 | output = model(data) 107 | test_loss += F.cross_entropy(output, target).data[0] 108 | pred = output.data.max(1)[1] # get the index of the max log-probability 109 | correct += pred.cpu().eq(indx_target).sum() 110 | 111 | test_loss = test_loss / len(test_loader) # average over number of mini-batch 112 | acc = 100. * correct / len(test_loader.dataset) 113 | print('\tTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format( 114 | test_loss, correct, len(test_loader.dataset), acc)) 115 | if acc > best_acc: 116 | new_file = os.path.join(args.logdir, 'best-{}.pth'.format(epoch)) 117 | misc.model_snapshot(model, new_file, old_file=old_file, verbose=True) 118 | best_acc = acc 119 | old_file = new_file 120 | except Exception as e: 121 | import traceback 122 | traceback.print_exc() 123 | finally: 124 | print("Total Elapse: {:.2f}, Best Result: {:.3f}%".format(time.time()-t_begin, best_acc)) 125 | 126 | 127 | -------------------------------------------------------------------------------- /svhn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aaron-xichen/pytorch-playground/5d476295f1c1564c8074eb885bb6138d6bd09889/svhn/__init__.py -------------------------------------------------------------------------------- /svhn/dataset.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torchvision import datasets, transforms 3 | from torch.utils.data import DataLoader 4 | import os 5 | 6 | def get(batch_size, data_root='/tmp/public_dataset/pytorch', train=True, val=True, **kwargs): 7 | data_root = os.path.expanduser(os.path.join(data_root, 'svhn-data')) 8 | num_workers = kwargs.setdefault('num_workers', 1) 9 | kwargs.pop('input_size', None) 10 | print("Building SVHN data loader with {} workers".format(num_workers)) 11 | 12 | def target_transform(target): 13 | return int(target) - 1 14 | 15 | ds = [] 16 | if train: 17 | train_loader = torch.utils.data.DataLoader( 18 | datasets.SVHN( 19 | root=data_root, split='train', download=True, 20 | transform=transforms.Compose([ 21 | transforms.ToTensor(), 22 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 23 | ]), 24 | target_transform=target_transform, 25 | ), 26 | batch_size=batch_size, shuffle=True, **kwargs) 27 | ds.append(train_loader) 28 | 29 | if val: 30 | test_loader = torch.utils.data.DataLoader( 31 | datasets.SVHN( 32 | root=data_root, split='test', download=True, 33 | transform=transforms.Compose([ 34 | transforms.ToTensor(), 35 | transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), 36 | ]), 37 | target_transform=target_transform 38 | ), 39 | batch_size=batch_size, shuffle=False, **kwargs) 40 | ds.append(test_loader) 41 | ds = ds[0] if len(ds) == 1 else ds 42 | return ds 43 | 44 | -------------------------------------------------------------------------------- /svhn/model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.utils.model_zoo as model_zoo 4 | import os 5 | from collections import OrderedDict 6 | from utee import misc 7 | print = misc.logger.info 8 | 9 | model_urls = { 10 | 'svhn': 'http://ml.cs.tsinghua.edu.cn/~chenxi/pytorch-models/svhn-f564f3d8.pth', 11 | } 12 | 13 | class SVHN(nn.Module): 14 | def __init__(self, features, n_channel, num_classes): 15 | super(SVHN, self).__init__() 16 | assert isinstance(features, nn.Sequential), type(features) 17 | self.features = features 18 | self.classifier = nn.Sequential( 19 | nn.Linear(n_channel, num_classes) 20 | ) 21 | print(self.features) 22 | print(self.classifier) 23 | 24 | def forward(self, x): 25 | x = self.features(x) 26 | x = x.view(x.size(0), -1) 27 | x = self.classifier(x) 28 | return x 29 | 30 | def make_layers(cfg, batch_norm=False): 31 | layers = [] 32 | in_channels = 3 33 | for i, v in enumerate(cfg): 34 | if v == 'M': 35 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 36 | else: 37 | padding = v[1] if isinstance(v, tuple) else 1 38 | out_channels = v[0] if isinstance(v, tuple) else v 39 | conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=padding) 40 | if batch_norm: 41 | layers += [conv2d, nn.BatchNorm2d(out_channels, affine=False), nn.ReLU(), nn.Dropout(0.3)] 42 | else: 43 | layers += [conv2d, nn.ReLU(), nn.Dropout(0.3)] 44 | in_channels = out_channels 45 | return nn.Sequential(*layers) 46 | 47 | def svhn(n_channel, pretrained=None): 48 | cfg = [n_channel, n_channel, 'M', 2*n_channel, 2*n_channel, 'M', 4*n_channel, 4*n_channel, 'M', (8*n_channel, 0), 'M'] 49 | layers = make_layers(cfg, batch_norm=True) 50 | model = SVHN(layers, n_channel=8*n_channel, num_classes=10) 51 | if pretrained is not None: 52 | m = model_zoo.load_url(model_urls['svhn']) 53 | state_dict = m.state_dict() if isinstance(m, nn.Module) else m 54 | assert isinstance(state_dict, (dict, OrderedDict)), type(state_dict) 55 | model.load_state_dict(state_dict) 56 | return model 57 | 58 | 59 | -------------------------------------------------------------------------------- /svhn/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import time 4 | 5 | from utee import misc 6 | import torch 7 | import torch.nn.functional as F 8 | import torch.optim as optim 9 | from torch.autograd import Variable 10 | 11 | import dataset 12 | import model 13 | from IPython import embed 14 | 15 | parser = argparse.ArgumentParser(description='PyTorch SVHN Example') 16 | parser.add_argument('--channel', type=int, default=32, help='first conv channel (default: 32)') 17 | parser.add_argument('--wd', type=float, default=0.001, help='weight decay') 18 | parser.add_argument('--batch_size', type=int, default=200, help='input batch size for training (default: 64)') 19 | parser.add_argument('--epochs', type=int, default=150, help='number of epochs to train (default: 10)') 20 | parser.add_argument('--lr', type=float, default=0.001, help='learning rate (default: 1e-3)') 21 | parser.add_argument('--gpu', default=None, help='index of gpus to use') 22 | parser.add_argument('--ngpu', type=int, default=2, help='number of gpus to use') 23 | parser.add_argument('--seed', type=int, default=117, help='random seed (default: 1)') 24 | parser.add_argument('--log_interval', type=int, default=100, help='how many batches to wait before logging training status') 25 | parser.add_argument('--test_interval', type=int, default=5, help='how many epochs to wait before another test') 26 | parser.add_argument('--logdir', default='log/default', help='folder to save to the log') 27 | parser.add_argument('--data_root', default='/tmp/public_dataset/pytorch/', help='folder to save the model') 28 | parser.add_argument('--decreasing_lr', default='80,120', help='decreasing strategy') 29 | args = parser.parse_args() 30 | args.logdir = os.path.join(os.path.dirname(__file__), args.logdir) 31 | misc.logger.init(args.logdir, 'train_log') 32 | print = misc.logger.info 33 | 34 | # select gpu 35 | args.gpu = misc.auto_select_gpu(utility_bound=0, num_gpu=args.ngpu, selected_gpus=args.gpu) 36 | args.ngpu = len(args.gpu) 37 | 38 | # logger 39 | misc.ensure_dir(args.logdir) 40 | print("=================FLAGS==================") 41 | for k, v in args.__dict__.items(): 42 | print('{}: {}'.format(k, v)) 43 | print("========================================") 44 | 45 | # seed 46 | args.cuda = torch.cuda.is_available() 47 | torch.manual_seed(args.seed) 48 | if args.cuda: 49 | torch.cuda.manual_seed(args.seed) 50 | 51 | # data loader and model 52 | train_loader, test_loader = dataset.get(batch_size=args.batch_size, data_root=args.data_root, num_workers=1) 53 | model = model.svhn(n_channel=args.channel) 54 | model = torch.nn.DataParallel(model, device_ids= range(args.ngpu)) 55 | if args.cuda: 56 | model.cuda() 57 | 58 | # optimizer 59 | optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd) 60 | decreasing_lr = list(map(int, args.decreasing_lr.split(','))) 61 | print('decreasing_lr: ' + str(decreasing_lr)) 62 | best_acc, old_file = 0, None 63 | t_begin = time.time() 64 | try: 65 | for epoch in range(args.epochs): 66 | model.train() 67 | if epoch in decreasing_lr: 68 | optimizer.param_groups[0]['lr'] *= 0.1 69 | for batch_idx, (data, target) in enumerate(train_loader): 70 | indx_target = target.clone() 71 | if args.cuda: 72 | data, target = data.cuda(), target.cuda() 73 | data, target = Variable(data), Variable(target) 74 | 75 | optimizer.zero_grad() 76 | output = model(data) 77 | loss = F.cross_entropy(output, target) 78 | loss.backward() 79 | optimizer.step() 80 | 81 | if batch_idx % args.log_interval == 0 and batch_idx > 0: 82 | pred = output.data.max(1)[1] # get the index of the max log-probability 83 | correct = pred.cpu().eq(indx_target).sum() 84 | acc = correct * 1.0 / len(data) 85 | print('Train Epoch: {} [{}/{}] Loss: {:.6f} Acc: {:.4f} lr: {:.2e}'.format( 86 | epoch, batch_idx * len(data), len(train_loader.dataset), 87 | loss.data[0], acc, optimizer.param_groups[0]['lr'])) 88 | 89 | elapse_time = time.time() - t_begin 90 | speed_epoch = elapse_time / (epoch + 1) 91 | speed_batch = speed_epoch / len(train_loader) 92 | eta = speed_epoch * args.epochs - elapse_time 93 | print("Elapsed {:.2f}s, {:.2f} s/epoch, {:.2f} s/batch, ets {:.2f}s".format( 94 | elapse_time, speed_epoch, speed_batch, eta)) 95 | misc.model_snapshot(model, os.path.join(args.logdir, 'latest.pth')) 96 | 97 | if epoch % args.test_interval == 0: 98 | model.eval() 99 | test_loss = 0 100 | correct = 0 101 | for data, target in test_loader: 102 | indx_target = target.clone() 103 | if args.cuda: 104 | data, target = data.cuda(), target.cuda().long().squeeze() 105 | data, target = Variable(data, volatile=True), Variable(target) 106 | output = model(data) 107 | test_loss += F.cross_entropy(output, target).data[0] 108 | pred = output.data.max(1)[1] # get the index of the max log-probability 109 | correct += pred.cpu().eq(indx_target).sum() 110 | 111 | test_loss = test_loss / len(test_loader) # average over number of mini-batch 112 | acc = 100. * correct / len(test_loader.dataset) 113 | print('\tTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format( 114 | test_loss, correct, len(test_loader.dataset), acc)) 115 | if acc > best_acc: 116 | new_file = os.path.join(args.logdir, 'best-{}.pth'.format(epoch)) 117 | misc.model_snapshot(model, new_file, old_file=old_file, verbose=True) 118 | best_acc = acc 119 | old_file = new_file 120 | except Exception as e: 121 | import traceback 122 | traceback.print_exc() 123 | finally: 124 | print("Total Elapse: {:.2f}, Best Result: {:.3f}%".format(time.time()-t_begin, best_acc)) 125 | 126 | 127 | -------------------------------------------------------------------------------- /utee/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aaron-xichen/pytorch-playground/5d476295f1c1564c8074eb885bb6138d6bd09889/utee/__init__.py -------------------------------------------------------------------------------- /utee/misc.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import os 3 | import shutil 4 | import pickle as pkl 5 | import time 6 | import numpy as np 7 | import hashlib 8 | 9 | from IPython import embed 10 | 11 | class Logger(object): 12 | def __init__(self): 13 | self._logger = None 14 | 15 | def init(self, logdir, name='log'): 16 | if self._logger is None: 17 | import logging 18 | if not os.path.exists(logdir): 19 | os.makedirs(logdir) 20 | log_file = os.path.join(logdir, name) 21 | if os.path.exists(log_file): 22 | os.remove(log_file) 23 | self._logger = logging.getLogger() 24 | self._logger.setLevel('INFO') 25 | fh = logging.FileHandler(log_file) 26 | ch = logging.StreamHandler() 27 | self._logger.addHandler(fh) 28 | self._logger.addHandler(ch) 29 | 30 | def info(self, str_info): 31 | self.init('/tmp', 'tmp.log') 32 | self._logger.info(str_info) 33 | logger = Logger() 34 | 35 | print = logger.info 36 | def ensure_dir(path, erase=False): 37 | if os.path.exists(path) and erase: 38 | print("Removing old folder {}".format(path)) 39 | shutil.rmtree(path) 40 | if not os.path.exists(path): 41 | print("Creating folder {}".format(path)) 42 | os.makedirs(path) 43 | 44 | def load_pickle(path): 45 | begin_st = time.time() 46 | with open(path, 'rb') as f: 47 | print("Loading pickle object from {}".format(path)) 48 | v = pkl.load(f) 49 | print("=> Done ({:.4f} s)".format(time.time() - begin_st)) 50 | return v 51 | 52 | def dump_pickle(obj, path): 53 | with open(path, 'wb') as f: 54 | print("Dumping pickle object to {}".format(path)) 55 | pkl.dump(obj, f, protocol=pkl.HIGHEST_PROTOCOL) 56 | 57 | def auto_select_gpu(mem_bound=500, utility_bound=0, gpus=(0, 1, 2, 3, 4, 5, 6, 7), num_gpu=1, selected_gpus=None): 58 | import sys 59 | import os 60 | import subprocess 61 | import re 62 | import time 63 | import numpy as np 64 | if 'CUDA_VISIBLE_DEVCIES' in os.environ: 65 | sys.exit(0) 66 | if selected_gpus is None: 67 | mem_trace = [] 68 | utility_trace = [] 69 | for i in range(5): # sample 5 times 70 | info = subprocess.check_output('nvidia-smi', shell=True).decode('utf-8') 71 | mem = [int(s[:-5]) for s in re.compile('\d+MiB\s/').findall(info)] 72 | utility = [int(re.compile('\d+').findall(s)[0]) for s in re.compile('\d+%\s+Default').findall(info)] 73 | mem_trace.append(mem) 74 | utility_trace.append(utility) 75 | time.sleep(0.1) 76 | mem = np.mean(mem_trace, axis=0) 77 | utility = np.mean(utility_trace, axis=0) 78 | assert(len(mem) == len(utility)) 79 | nGPU = len(utility) 80 | ideal_gpus = [i for i in range(nGPU) if mem[i] <= mem_bound and utility[i] <= utility_bound and i in gpus] 81 | 82 | if len(ideal_gpus) < num_gpu: 83 | print("No sufficient resource, available: {}, require {} gpu".format(ideal_gpus, num_gpu)) 84 | sys.exit(0) 85 | else: 86 | selected_gpus = list(map(str, ideal_gpus[:num_gpu])) 87 | else: 88 | selected_gpus = selected_gpus.split(',') 89 | 90 | print("Setting GPU: {}".format(selected_gpus)) 91 | os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(selected_gpus) 92 | return selected_gpus 93 | 94 | def expand_user(path): 95 | return os.path.abspath(os.path.expanduser(path)) 96 | 97 | def model_snapshot(model, new_file, old_file=None, verbose=False): 98 | from collections import OrderedDict 99 | import torch 100 | if isinstance(model, torch.nn.DataParallel): 101 | model = model.module 102 | if old_file and os.path.exists(expand_user(old_file)): 103 | if verbose: 104 | print("Removing old model {}".format(expand_user(old_file))) 105 | os.remove(expand_user(old_file)) 106 | if verbose: 107 | print("Saving model to {}".format(expand_user(new_file))) 108 | 109 | state_dict = OrderedDict() 110 | for k, v in model.state_dict().items(): 111 | if v.is_cuda: 112 | v = v.cpu() 113 | state_dict[k] = v 114 | torch.save(state_dict, expand_user(new_file)) 115 | 116 | 117 | def load_lmdb(lmdb_file, n_records=None): 118 | import lmdb 119 | import numpy as np 120 | lmdb_file = expand_user(lmdb_file) 121 | if os.path.exists(lmdb_file): 122 | data = [] 123 | env = lmdb.open(lmdb_file, readonly=True, max_readers=512) 124 | with env.begin() as txn: 125 | cursor = txn.cursor() 126 | begin_st = time.time() 127 | print("Loading lmdb file {} into memory".format(lmdb_file)) 128 | for key, value in cursor: 129 | _, target, _ = key.decode('ascii').split(':') 130 | target = int(target) 131 | img = cv2.imdecode(np.fromstring(value, np.uint8), cv2.IMREAD_COLOR) 132 | data.append((img, target)) 133 | if n_records is not None and len(data) >= n_records: 134 | break 135 | env.close() 136 | print("=> Done ({:.4f} s)".format(time.time() - begin_st)) 137 | return data 138 | else: 139 | print("Not found lmdb file".format(lmdb_file)) 140 | 141 | def str2img(str_b): 142 | return cv2.imdecode(np.fromstring(str_b, np.uint8), cv2.IMREAD_COLOR) 143 | 144 | def img2str(img): 145 | return cv2.imencode('.jpg', img)[1].tostring() 146 | 147 | def md5(s): 148 | m = hashlib.md5() 149 | m.update(s) 150 | return m.hexdigest() 151 | 152 | def eval_model(model, ds, n_sample=None, ngpu=1, is_imagenet=False): 153 | import tqdm 154 | import torch 155 | from torch import nn 156 | from torch.autograd import Variable 157 | 158 | class ModelWrapper(nn.Module): 159 | def __init__(self, model): 160 | super(ModelWrapper, self).__init__() 161 | self.model = model 162 | self.mean = [0.485, 0.456, 0.406] 163 | self.std = [0.229, 0.224, 0.225] 164 | 165 | def forward(self, input): 166 | input.data.div_(255.) 167 | input.data[:, 0, :, :].sub_(self.mean[0]).div_(self.std[0]) 168 | input.data[:, 1, :, :].sub_(self.mean[1]).div_(self.std[1]) 169 | input.data[:, 2, :, :].sub_(self.mean[2]).div_(self.std[2]) 170 | return self.model(input) 171 | 172 | correct1, correct5 = 0, 0 173 | n_passed = 0 174 | if is_imagenet: 175 | model = ModelWrapper(model) 176 | model = model.eval() 177 | model = torch.nn.DataParallel(model, device_ids=range(ngpu)).cuda() 178 | 179 | n_sample = len(ds) if n_sample is None else n_sample 180 | for idx, (data, target) in enumerate(tqdm.tqdm(ds, total=n_sample)): 181 | n_passed += len(data) 182 | data = Variable(torch.FloatTensor(data)).cuda() 183 | indx_target = torch.LongTensor(target) 184 | output = model(data) 185 | bs = output.size(0) 186 | idx_pred = output.data.sort(1, descending=True)[1] 187 | 188 | idx_gt1 = indx_target.expand(1, bs).transpose_(0, 1) 189 | idx_gt5 = idx_gt1.expand(bs, 5) 190 | 191 | correct1 += idx_pred[:, :1].cpu().eq(idx_gt1).sum() 192 | correct5 += idx_pred[:, :5].cpu().eq(idx_gt5).sum() 193 | 194 | if idx >= n_sample - 1: 195 | break 196 | 197 | acc1 = correct1 * 1.0 / n_passed 198 | acc5 = correct5 * 1.0 / n_passed 199 | return acc1, acc5 200 | 201 | def load_state_dict(model, model_urls, model_root): 202 | from torch.utils import model_zoo 203 | from torch import nn 204 | import re 205 | from collections import OrderedDict 206 | own_state_old = model.state_dict() 207 | own_state = OrderedDict() # remove all 'group' string 208 | for k, v in own_state_old.items(): 209 | k = re.sub('group\d+\.', '', k) 210 | own_state[k] = v 211 | 212 | state_dict = model_zoo.load_url(model_urls, model_root) 213 | 214 | for name, param in state_dict.items(): 215 | if name not in own_state: 216 | print(own_state.keys()) 217 | raise KeyError('unexpected key "{}" in state_dict' 218 | .format(name)) 219 | if isinstance(param, nn.Parameter): 220 | # backwards compatibility for serialized parameters 221 | param = param.data 222 | own_state[name].copy_(param) 223 | 224 | missing = set(own_state.keys()) - set(state_dict.keys()) 225 | no_use = set(state_dict.keys()) - set(own_state.keys()) 226 | if len(no_use) > 0: 227 | raise KeyError('some keys are not used: "{}"'.format(no_use)) 228 | 229 | -------------------------------------------------------------------------------- /utee/quant.py: -------------------------------------------------------------------------------- 1 | from torch.autograd import Variable 2 | import torch 3 | from torch import nn 4 | from collections import OrderedDict 5 | import math 6 | from IPython import embed 7 | 8 | def compute_integral_part(input, overflow_rate): 9 | abs_value = input.abs().view(-1) 10 | sorted_value = abs_value.sort(dim=0, descending=True)[0] 11 | split_idx = int(overflow_rate * len(sorted_value)) 12 | v = sorted_value[split_idx] 13 | if isinstance(v, Variable): 14 | v = float(v.data.cpu()) 15 | sf = math.ceil(math.log2(v+1e-12)) 16 | return sf 17 | 18 | def linear_quantize(input, sf, bits): 19 | assert bits >= 1, bits 20 | if bits == 1: 21 | return torch.sign(input) - 1 22 | delta = math.pow(2.0, -sf) 23 | bound = math.pow(2.0, bits-1) 24 | min_val = - bound 25 | max_val = bound - 1 26 | rounded = torch.floor(input / delta + 0.5) 27 | 28 | clipped_value = torch.clamp(rounded, min_val, max_val) * delta 29 | return clipped_value 30 | 31 | def log_minmax_quantize(input, bits): 32 | assert bits >= 1, bits 33 | if bits == 1: 34 | return torch.sign(input), 0.0, 0.0 35 | 36 | s = torch.sign(input) 37 | input0 = torch.log(torch.abs(input) + 1e-20) 38 | v = min_max_quantize(input0, bits-1) 39 | v = torch.exp(v) * s 40 | return v 41 | 42 | def log_linear_quantize(input, sf, bits): 43 | assert bits >= 1, bits 44 | if bits == 1: 45 | return torch.sign(input), 0.0, 0.0 46 | 47 | s = torch.sign(input) 48 | input0 = torch.log(torch.abs(input) + 1e-20) 49 | v = linear_quantize(input0, sf, bits-1) 50 | v = torch.exp(v) * s 51 | return v 52 | 53 | def min_max_quantize(input, bits): 54 | assert bits >= 1, bits 55 | if bits == 1: 56 | return torch.sign(input) - 1 57 | min_val, max_val = input.min(), input.max() 58 | 59 | if isinstance(min_val, Variable): 60 | max_val = float(max_val.data.cpu().numpy()[0]) 61 | min_val = float(min_val.data.cpu().numpy()[0]) 62 | 63 | input_rescale = (input - min_val) / (max_val - min_val) 64 | 65 | n = math.pow(2.0, bits) - 1 66 | v = torch.floor(input_rescale * n + 0.5) / n 67 | 68 | v = v * (max_val - min_val) + min_val 69 | return v 70 | 71 | def tanh_quantize(input, bits): 72 | assert bits >= 1, bits 73 | if bits == 1: 74 | return torch.sign(input) 75 | input = torch.tanh(input) # [-1, 1] 76 | input_rescale = (input + 1.0) / 2 #[0, 1] 77 | n = math.pow(2.0, bits) - 1 78 | v = torch.floor(input_rescale * n + 0.5) / n 79 | v = 2 * v - 1 # [-1, 1] 80 | 81 | v = 0.5 * torch.log((1 + v) / (1 - v)) # arctanh 82 | return v 83 | 84 | 85 | class LinearQuant(nn.Module): 86 | def __init__(self, name, bits, sf=None, overflow_rate=0.0, counter=10): 87 | super(LinearQuant, self).__init__() 88 | self.name = name 89 | self._counter = counter 90 | 91 | self.bits = bits 92 | self.sf = sf 93 | self.overflow_rate = overflow_rate 94 | 95 | @property 96 | def counter(self): 97 | return self._counter 98 | 99 | def forward(self, input): 100 | if self._counter > 0: 101 | self._counter -= 1 102 | sf_new = self.bits - 1 - compute_integral_part(input, self.overflow_rate) 103 | self.sf = min(self.sf, sf_new) if self.sf is not None else sf_new 104 | return input 105 | else: 106 | output = linear_quantize(input, self.sf, self.bits) 107 | return output 108 | 109 | def __repr__(self): 110 | return '{}(sf={}, bits={}, overflow_rate={:.3f}, counter={})'.format( 111 | self.__class__.__name__, self.sf, self.bits, self.overflow_rate, self.counter) 112 | 113 | class LogQuant(nn.Module): 114 | def __init__(self, name, bits, sf=None, overflow_rate=0.0, counter=10): 115 | super(LogQuant, self).__init__() 116 | self.name = name 117 | self._counter = counter 118 | 119 | self.bits = bits 120 | self.sf = sf 121 | self.overflow_rate = overflow_rate 122 | 123 | @property 124 | def counter(self): 125 | return self._counter 126 | 127 | def forward(self, input): 128 | if self._counter > 0: 129 | self._counter -= 1 130 | log_abs_input = torch.log(torch.abs(input)) 131 | sf_new = self.bits - 1 - compute_integral_part(log_abs_input, self.overflow_rate) 132 | self.sf = min(self.sf, sf_new) if self.sf is not None else sf_new 133 | return input 134 | else: 135 | output = log_linear_quantize(input, self.sf, self.bits) 136 | return output 137 | 138 | def __repr__(self): 139 | return '{}(sf={}, bits={}, overflow_rate={:.3f}, counter={})'.format( 140 | self.__class__.__name__, self.sf, self.bits, self.overflow_rate, self.counter) 141 | 142 | class NormalQuant(nn.Module): 143 | def __init__(self, name, bits, quant_func): 144 | super(NormalQuant, self).__init__() 145 | self.name = name 146 | self.bits = bits 147 | self.quant_func = quant_func 148 | 149 | @property 150 | def counter(self): 151 | return self._counter 152 | 153 | def forward(self, input): 154 | output = self.quant_func(input, self.bits) 155 | return output 156 | 157 | def __repr__(self): 158 | return '{}(bits={})'.format(self.__class__.__name__, self.bits) 159 | 160 | def duplicate_model_with_quant(model, bits, overflow_rate=0.0, counter=10, type='linear'): 161 | """assume that original model has at least a nn.Sequential""" 162 | assert type in ['linear', 'minmax', 'log', 'tanh'] 163 | if isinstance(model, nn.Sequential): 164 | l = OrderedDict() 165 | for k, v in model._modules.items(): 166 | if isinstance(v, (nn.Conv2d, nn.Linear, nn.BatchNorm1d, nn.BatchNorm2d, nn.AvgPool2d)): 167 | l[k] = v 168 | if type == 'linear': 169 | quant_layer = LinearQuant('{}_quant'.format(k), bits=bits, overflow_rate=overflow_rate, counter=counter) 170 | elif type == 'log': 171 | # quant_layer = LogQuant('{}_quant'.format(k), bits=bits, overflow_rate=overflow_rate, counter=counter) 172 | quant_layer = NormalQuant('{}_quant'.format(k), bits=bits, quant_func=log_minmax_quantize) 173 | elif type == 'minmax': 174 | quant_layer = NormalQuant('{}_quant'.format(k), bits=bits, quant_func=min_max_quantize) 175 | else: 176 | quant_layer = NormalQuant('{}_quant'.format(k), bits=bits, quant_func=tanh_quantize) 177 | l['{}_{}_quant'.format(k, type)] = quant_layer 178 | else: 179 | l[k] = duplicate_model_with_quant(v, bits, overflow_rate, counter, type) 180 | m = nn.Sequential(l) 181 | return m 182 | else: 183 | for k, v in model._modules.items(): 184 | model._modules[k] = duplicate_model_with_quant(v, bits, overflow_rate, counter, type) 185 | return model 186 | 187 | -------------------------------------------------------------------------------- /utee/selector.py: -------------------------------------------------------------------------------- 1 | from utee import misc 2 | import os 3 | from imagenet import dataset 4 | print = misc.logger.info 5 | from IPython import embed 6 | 7 | known_models = [ 8 | 'mnist', 'svhn', # 28x28 9 | 'cifar10', 'cifar100', # 32x32 10 | 'stl10', # 96x96 11 | 'alexnet', # 224x224 12 | 'vgg16', 'vgg16_bn', 'vgg19', 'vgg19_bn', # 224x224 13 | 'resnet18', 'resnet34', 'resnet50', 'resnet101','resnet152', # 224x224 14 | 'squeezenet_v0', 'squeezenet_v1', #224x224 15 | 'inception_v3', # 299x299 16 | ] 17 | 18 | def mnist(cuda=True, model_root=None): 19 | print("Building and initializing mnist parameters") 20 | from mnist import model, dataset 21 | m = model.mnist(pretrained=os.path.join(model_root, 'mnist.pth')) 22 | if cuda: 23 | m = m.cuda() 24 | return m, dataset.get, False 25 | 26 | def svhn(cuda=True, model_root=None): 27 | print("Building and initializing svhn parameters") 28 | from svhn import model, dataset 29 | m = model.svhn(32, pretrained=os.path.join(model_root, 'svhn.pth')) 30 | if cuda: 31 | m = m.cuda() 32 | return m, dataset.get, False 33 | 34 | def cifar10(cuda=True, model_root=None): 35 | print("Building and initializing cifar10 parameters") 36 | from cifar import model, dataset 37 | m = model.cifar10(128, pretrained=os.path.join(model_root, 'cifar10.pth')) 38 | if cuda: 39 | m = m.cuda() 40 | return m, dataset.get10, False 41 | 42 | def cifar100(cuda=True, model_root=None): 43 | print("Building and initializing cifar100 parameters") 44 | from cifar import model, dataset 45 | m = model.cifar100(128, pretrained=os.path.join(model_root, 'cifar100.pth')) 46 | if cuda: 47 | m = m.cuda() 48 | return m, dataset.get100, False 49 | 50 | def stl10(cuda=True, model_root=None): 51 | print("Building and initializing stl10 parameters") 52 | from stl10 import model, dataset 53 | m = model.stl10(32, pretrained=os.path.join(model_root, 'stl10.pth')) 54 | if cuda: 55 | m = m.cuda() 56 | return m, dataset.get, False 57 | 58 | def alexnet(cuda=True, model_root=None): 59 | print("Building and initializing alexnet parameters") 60 | from imagenet import alexnet as alx 61 | m = alx.alexnet(True, model_root) 62 | if cuda: 63 | m = m.cuda() 64 | return m, dataset.get, True 65 | 66 | def vgg16(cuda=True, model_root=None): 67 | print("Building and initializing vgg16 parameters") 68 | from imagenet import vgg 69 | m = vgg.vgg16(True, model_root) 70 | if cuda: 71 | m = m.cuda() 72 | return m, dataset.get, True 73 | 74 | def vgg16_bn(cuda=True, model_root=None): 75 | print("Building vgg16_bn parameters") 76 | from imagenet import vgg 77 | m = vgg.vgg16_bn(model_root) 78 | if cuda: 79 | m = m.cuda() 80 | return m, dataset.get, True 81 | 82 | def vgg19(cuda=True, model_root=None): 83 | print("Building and initializing vgg19 parameters") 84 | from imagenet import vgg 85 | m = vgg.vgg19(True, model_root) 86 | if cuda: 87 | m = m.cuda() 88 | return m, dataset.get, True 89 | 90 | def vgg19_bn(cuda=True, model_root=None): 91 | print("Building vgg19_bn parameters") 92 | from imagenet import vgg 93 | m = vgg.vgg19_bn(model_root) 94 | if cuda: 95 | m = m.cuda() 96 | return m, dataset.get, True 97 | 98 | def inception_v3(cuda=True, model_root=None): 99 | print("Building and initializing inception_v3 parameters") 100 | from imagenet import inception 101 | m = inception.inception_v3(True, model_root) 102 | if cuda: 103 | m = m.cuda() 104 | return m, dataset.get, True 105 | 106 | def resnet18(cuda=True, model_root=None): 107 | print("Building and initializing resnet-18 parameters") 108 | from imagenet import resnet 109 | m = resnet.resnet18(True, model_root) 110 | if cuda: 111 | m = m.cuda() 112 | return m, dataset.get, True 113 | 114 | def resnet34(cuda=True, model_root=None): 115 | print("Building and initializing resnet-34 parameters") 116 | from imagenet import resnet 117 | m = resnet.resnet34(True, model_root) 118 | if cuda: 119 | m = m.cuda() 120 | return m, dataset.get, True 121 | 122 | def resnet50(cuda=True, model_root=None): 123 | print("Building and initializing resnet-50 parameters") 124 | from imagenet import resnet 125 | m = resnet.resnet50(True, model_root) 126 | if cuda: 127 | m = m.cuda() 128 | return m, dataset.get, True 129 | 130 | def resnet101(cuda=True, model_root=None): 131 | print("Building and initializing resnet-101 parameters") 132 | from imagenet import resnet 133 | m = resnet.resnet101(True, model_root) 134 | if cuda: 135 | m = m.cuda() 136 | return m, dataset.get, True 137 | 138 | def resnet152(cuda=True, model_root=None): 139 | print("Building and initializing resnet-152 parameters") 140 | from imagenet import resnet 141 | m = resnet.resnet152(True, model_root) 142 | if cuda: 143 | m = m.cuda() 144 | return m, dataset.get, True 145 | 146 | def squeezenet_v0(cuda=True, model_root=None): 147 | print("Building and initializing squeezenet_v0 parameters") 148 | from imagenet import squeezenet 149 | m = squeezenet.squeezenet1_0(True, model_root) 150 | if cuda: 151 | m = m.cuda() 152 | return m, dataset.get, True 153 | 154 | def squeezenet_v1(cuda=True, model_root=None): 155 | print("Building and initializing squeezenet_v1 parameters") 156 | from imagenet import squeezenet 157 | m = squeezenet.squeezenet1_1(True, model_root) 158 | if cuda: 159 | m = m.cuda() 160 | return m, dataset.get, True 161 | 162 | def select(model_name, **kwargs): 163 | assert model_name in known_models, model_name 164 | kwargs.setdefault('model_root', os.path.expanduser('~/.torch/models')) 165 | return eval('{}'.format(model_name))(**kwargs) 166 | 167 | if __name__ == '__main__': 168 | m1 = alexnet() 169 | embed() 170 | 171 | 172 | --------------------------------------------------------------------------------