如何在 MNIST 训练的网络上测试我自己的图像？

Question

这是我第一次尝试训练网络并使用PyTorch，所以如果认为这很简单，请原谅我。

我有一个经过修改的预训练 AlexNet 网络，可以分类 3 类，我已经在 MNIST 上对其进行了训练，并将其映射到 3 个不同的标签。

class Net( nn.Module ) :
    def __init__( self ) :
        super( Net, self ).__init__()
        
        self.model = models.alexnet( pretrained = True )
        # changed in_channels from 3 to 1 bc images are black and white 
        self.model.features[ 0 ] = nn.Conv2d( 1, 64, kernel_size = 11, stride = 4, padding = 2 )
        
        # binary classifier -> 3 out_features
        self.model.classifier[ 4 ] = nn.Linear( 4096, 1024 )
        self.model.classifier[ 6 ] = nn.Linear( 1024, 3 )
        
    def forward( self, x ):
        return self.model( x )
    
model = Net().to( device )

我想在我绘制的单个 .png 图像上对此进行测试，该图像已经是 255x255，并且是黑白的。我想要预测的标签。这是我目前用于预处理图像的代码：

from PIL import Image
import matplotlib.pyplot as plt
import cv2
image_8 = Image.open( "eight.png" ).convert('L')

image_8 = list( image_8.getdata())
normalized_8 = [(255 - x) * 1.0 / 255.0 for x in image_8 ]
tensor_8 = torch.FloatTensor( normalized_8 )

pred = model( tensor_8 )

我从中得到以下错误：Expected 4-dimensional input for 4-dimensional weight [64, 1, 11, 11], but got 1-dimensional input of size [50176] instead。所以这显然是错误的做事方式，但我不确定如何进行。

Answer 1

将您的推理代码更改为以下内容。图像不打算被展平为 1d。

import matplotlib.pyplot as plt
import cv2
image_8 = cv2.imread("eight.png")

# following line may or may not be necessary
image_8 = cv2.cvtColor(image_8, cv2.COLOR_BGR2GRAY)

# you can divide numpy arrays by a constant natively
image_8 /= 255.

# This makes a 4d tensor (batched image) with shape [1, channels, width, height]
image_8 = torch.Tensor(tensor_8).unsqueeze(axis=0)

pred = model(image_8)

如果图像仍然是 3d（[1, width, height] 的形状），添加第二个 .unsqueeze(axis=0)。

如何在 MNIST 训练的网络上测试我自己的图像？

How to test my own image on a MNIST trained network?

neural-network

mnist

pytorch