PyTorch报错：仅支持空间目标批次(3D张量)，但目标维度为1的解决咨询

问题代码与报错 #

用户代码如下：

数据加载部分 #

data = ImageFolder(data_dir, transform=transforms.Compose([transforms.Resize((224,224)),transforms.ToTensor()]))

trainloader = torch.utils.data.DataLoader(data, batch_size=3600,
                                          shuffle=True, num_workers=2)

dataiter = iter(trainloader)
x_train, y_train = dataiter.next()

print(x_train.size())
print(y_train.size())

输出：

torch.Size([3600, 3, 224, 224])
torch.Size([3600])

模型与损失计算部分 #

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = torch.nn.Linear(224, 12)
        self.layer2 = torch.nn.Linear(12, 10)

    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = self.layer2(x)
        return x

net = Net()
y = net.forward(x_train)
lossFn = torch.nn.CrossEntropyLoss() 
loss = lossFn(y, y_train)
print(loss)

报错信息（翻译为中文）：

RuntimeError: 仅支持空间目标的批次（3D张量），但得到的目标维度为1

问题原因 #

1. 输入维度不匹配：输入的x_train是4D张量[批量数, 通道数, 高度, 宽度]，但第一个全连接层Linear(224,12)仅接收最后一维为224的张量，直接传入会导致输出维度混乱。
2. 未展平图像张量：全连接层需要一维特征向量，而非2D空间图像数据，必须先将4D图像张量展平为2D张量[批量数, 总特征数]。

修复方案 #

修改模型的__init__和forward方法，调整全连接层输入维度并添加张量展平操作：

import torch.nn.functional as F

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 输入维度为3通道×224高度×224宽度的总特征数
        self.layer1 = torch.nn.Linear(3 * 224 * 224, 12)
        self.layer2 = torch.nn.Linear(12, 10)

    def forward(self, x):
        # 展平张量：保留批量维度，将剩余维度合并为一维特征
        x = x.view(x.size(0), -1)
        x = F.relu(self.layer1(x))
        x = self.layer2(x)
        return x

net = Net()
# 直接调用模型即可，无需显式调用forward方法
y = net(x_train)
lossFn = torch.nn.CrossEntropyLoss() 
loss = lossFn(y, y_train)
print(loss)

修复说明 #

• 展平操作x.view(x.size(0), -1)将4D图像张量转为2D张量[3600, 3*224*224]，符合全连接层的输入要求。
• 调整第一个全连接层的输入维度为3*224*224，与展平后的特征数匹配，最终模型输出为[3600,10]，和目标张量y_train的[3600]维度匹配，CrossEntropyLoss可正常计算损失。

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