pytorch-playground

这是一个面向 PyTorch 初学者的游乐场，其中包含在流行数据集上预定义的模型。目前我们支持：

• mnist、svhn
• cifar10、cifar100
• stl10
• alexnet
• vgg16、vgg16_bn、vgg19、vgg19_bn
• resnet18、resnet34、resnet50、resnet101、resnet152
• squeezenet_v0、squeezenet_v1
• inception_v3

以下是 MNIST 数据集的示例。这将自动下载数据集和预训练模型。

import torch
from torch.autograd import Variable
from utee import selector
model_raw, ds_fetcher, is_imagenet = selector.select('mnist')
ds_val = ds_fetcher(batch_size=10, train=False, val=True)
for idx, (data, target) in enumerate(ds_val):
    data =  Variable(torch.FloatTensor(data)).cuda()
    output = model_raw(data)

此外，如果要在 mnist 上训练 MLP 模型，只需运行 Python mnist/train.py

安装

python3 setup.py develop --user

ImageNet 数据集

我们提供了预计算好的 ImageNet 验证数据集，尺寸为 224x224x3。我们首先将图像的短边缩放到 256，然后在中心裁剪出 224x224 的图像。接着，我们将裁剪后的图像编码为 jpg 字符串，并转储到 pickle 文件中。

• cd script
• 下载 val224_compressed.pkl（清华镜像 / Google Drive）
• Python convert.py（需要 48G 内存，感谢 @jnorwood）

量化

我们还提供了一个简单的演示，用于将这些模型量化到指定的位宽，支持多种方法，包括线性方法、最小最大值方法和非线性方法。

quantize --type cifar10 --quant_method linear --param_bits 8 --fwd_bits 8 --bn_bits 8 --ngpu 1

顶级准确率

我们使用线性量化方法评估了流行数据集和模型的性能。所有结果中，BN中的运行均值和运行方差的位宽均为10位（32位浮点除外）。