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

Question

I want to train SqueezeNet 1.1 model using MNIST dataset instead of ImageNet dataset.
Can i have the same model as torchvision.models.squeezenet?
Thanks!

Answer 1

TorchVision provides only ImageNet data pretrained model for the SqueezeNet architecture. However, you can train your own model using MNIST dataset by taking only the model (but not the pre-trained one) from 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

Now, you can use this architecture to train a model on MNIST data, which should not take too long.

---

One modification to keep in mind is to update the number of classes which is 10 for MNIST. Specifically, the 1000 should be changed to 10, and the kernel and stride accordingly.

  (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)
  )

Here's the relevant explanation: finetuning_torchvision_models-squeezenet

Answer 2

The initialization of the pretrained weights is possible but you'll get trouble with the strides and kernel sizes since MNIST images are 28X28 pixels. Most probably the reduction will lead to (batch_sizex1x1xchannel) feature maps before the net is at its infernece layer which will then cause an error.