PyTorch模型训练正常但单张同尺寸图片测试时矩阵相乘报错

基于64x64x3的数据集训练了PyTorch模型，训练和评估过程均正常，但测试实时单张图片时，尽管输入尺寸显示为torch.Size([3, 64, 64])，却触发RuntimeError: mat1 and mat2 shapes cannot be multiplied错误。

模型代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F

class liveClassifier(nn.Module):
  def __init__(self):
    super(liveClassifier, self).__init__()
    # input: 64x64x3
    self.conv01 = nn.Conv2d(in_channels=3, out_channels=8, kernel_size=3)
    # => 62x62x8
    self.conv02 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3)
    # => 60x60x16
    self.pool01 = nn.MaxPool2d(kernel_size=2)
    # => 30x30x16
    self.conv03 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3)
    # => 28x28x32
    self.pool02 = nn.MaxPool2d(kernel_size=2)
    # => 14x14x32
    self.conv04 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
    # => 12x12x64
    self.pool03 = nn.MaxPool2d(kernel_size=2)
    # => 6x6x64
    self.linear01 = nn.Linear(6 * 6 * 64, 32)
    self.linear02 = nn.Linear(32, 16)
    self.linear03 = nn.Linear(16, 3)

  def forward(self, input):
    x = self.conv01(input)
    x = F.relu(x)
    
    x = self.conv02(x)
    x = F.relu(x)

    x = self.pool01(x)

    x = self.conv03(x)
    x = F.relu(x)

    x = self.pool02(x)

    x = self.conv04(x)
    x = F.relu(x)

    x = self.pool03(x)

    x = x.flatten(start_dim=1)
    
    x = self.linear01(x)
    x = F.relu(x)

    x = self.linear02(x)
    x = F.relu(x)

    x = self.linear03(x)
    
    output = F.softmax(x, dim=1)
    return output

训练代码 #

num_epochs = 30

train_loss_list = []
train_acc_list = []

validation_loss_list = []
validation_acc_list = []

for epoch in range(num_epochs):
  
  train_running_loss = 0.0
  train_acc = 0.0

  model = model.train()

  for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)

        ## forward + backprop + loss
        logits = model(images)
        loss = criterion(logits, labels)
        optimizer.zero_grad()
        loss.backward()

        ## update model params
        optimizer.step()

        train_running_loss += loss.detach().item()
        train_acc += get_accuracy(logits, labels, BATCH_SIZE)
    

  
  model = model.eval()
  validation_running_loss = 0.0
  validation_acc = 0.0
  for i, (images, labels) in enumerate(validation_loader, 0):
    images = images.to(device)
    labels = labels.to(device)
    outputs = model(images)
    loss = criterion(outputs, labels)
    validation_running_loss += loss.detach().item()
    validation_acc += get_accuracy(outputs, labels, BATCH_SIZE)
  
  train_loss_list.append(train_running_loss / i)
  train_acc_list.append(train_acc / i)
  validation_loss_list.append(validation_running_loss / i)
  validation_acc_list.append(validation_acc / i)
  
  print('Epoch: %d | Train Loss: %.4f | Train Accuracy: %.2f' \
          %(epoch, train_running_loss / i, train_acc / i))
  print('Validation Loss: %.4f | Validation Accuracy: %.2f' \
          %(validation_running_loss / i, validation_acc / i))


model = model.eval()

测试代码（触发错误） #

print(type(frame))
frame = transform(Image.fromarray(frame)).float().to(device)
print(frame.shape) # torch.Size([3, 64, 64])
model.eval()
print(model(frame))

训练时train_loader输出的是批量张量，形状为[BATCH_SIZE, 3, 64, 64]，而测试时传入的是单张图片张量，缺少批量维度，形状为[3, 64, 64]。

模型的卷积层可以兼容任意维度（只要通道、尺寸匹配），但进入全连接层时，x.flatten(start_dim=1)对单张图片会将张量展平为[6*6*64=2304]，而全连接层linear01的权重形状为[32, 2304]。矩阵乘法要求第一个张量的最后一维等于第二个张量的倒数第二维，此时mat1是[2304]，mat2是[32,2304]，维度不匹配，因此触发乘法错误。

给单张图片张量添加批量维度即可，有两种常用方式：

方式1：修改测试代码添加批量维度 #

print(type(frame))
frame = transform(Image.fromarray(frame)).float().to(device)
# 添加第0维，形状变为[1, 3, 64, 64]
frame = frame.unsqueeze(0)
model.eval()
print(model(frame))

方式2：在模型中自动处理输入维度 #

修改模型的forward方法，自动检查并添加批量维度：

def forward(self, input):
    # 若输入无批量维度，自动添加第0维
    if len(input.shape) == 3:
        input = input.unsqueeze(0)
    x = self.conv01(input)
    x = F.relu(x)
    
    x = self.conv02(x)
    x = F.relu(x)

    x = self.pool01(x)

    x = self.conv03(x)
    x = F.relu(x)

    x = self.pool02(x)

    x = self.conv04(x)
    x = F.relu(x)

    x = self.pool03(x)

    x = x.flatten(start_dim=1)
    
    x = self.linear01(x)
    x = F.relu(x)

    x = self.linear02(x)
    x = F.relu(x)

    x = self.linear03(x)
    
    output = F.softmax(x, dim=1)
    return output

内容的提问来源于stack exchange，提问作者Shy Cohen