使用 MNIST 数据集 Pytorch 训练 SqueezeNet 模型

Train SqueezeNet model using MNIST dataset Pytorch

我想使用 MNIST 数据集而不是 ImageNet 数据集来训练 SqueezeNet 1.1 模型。
我可以拥有与torchvision.models.squeezenet相同的型号吗?
谢谢!

TorchVision 只为 SqueezeNet 架构提供 ImageNet 数据预训练模型。但是,您可以使用 MNIST 数据集训练您自己的模型,只需使用 torchvision.models.

中的模型(而不是预训练模型)
In [10]: import torchvision as tv

# get the model architecture only; ignore `pretrained` flag
In [11]: squeezenet11 = tv.models.squeezenet1_1()

In [12]: squeezenet11.training   
Out[12]: True

现在,您可以使用此架构在 MNIST 数据上训练模型,这不会花费太长时间。


要记住的一个修改是更新 类 的数量,对于 MNIST 为 10。具体来说,应该把1000改成10,kernel和stride也相应的改。

  (classifier): Sequential(
    (0): Dropout(p=0.5)
    (1): Conv2d(512, 1000, kernel_size=(1, 1), stride=(1, 1))
    (2): ReLU(inplace)
    (3): AvgPool2d(kernel_size=13, stride=1, padding=0)
  )

相关解释如下:finetuning_torchvision_models-squeezenet

预训练权重的初始化是可能的,但由于 MNIST 图像是 28X28 像素,因此您会遇到步幅和内核大小方面的问题。最有可能的是,在网络处于其推理层之前,减少将导致 (batch_sizex1x1xchannel) 个特征映射,这将导致错误。