如何调整一批图像的大小以用于Pytorch线性回归

Question

我正在尝试创建一个简单的线性回归神经网络，用于批量图像。输入维度是[BatchSize, 3, Width, Height]，第二维度表示输入图像的RGB通道。

以下是我对回归模型的(失败的(尝试：

class LinearNet(torch.nn.Module):
def __init__(self, Chn, W,H, nHidden):
"""
Input: A [BatchSize x Channels x Width x Height] set of images
Output: A fitted regression model with weights dimension : [Width x Height]
"""
super(LinearNet, self).__init__()
self.Chn = Chn
self.W = W
self.H = H
self.hidden = torch.nn.Linear(Chn*W*H,nHidden)   # hidden layer
self.predict = torch.nn.Linear(nHidden, Chn*W*H)   # output layer
def forward(self, x):
torch.reshape(x, (-1,self.Chn*self.W*self.H)) # FAILS here
# x = x.resize(-1,self.Chn*self.W*self.H)  
x = F.relu(self.hidden(x))      # activation function for hidden layer
x = self.predict(x)             # linear output
x = x.resize(-1,self.Chn, self.W,self.H)
return x

当发送一批尺寸为[128 x 3 x 96 x 128]的图像时，这在指示行上失败：

RuntimeError: mat1 and mat2 shapes cannot be multiplied (36864x128 and 36864x256)

应该如何正确操作矩阵维度以使用这些pytorch函数？

更新基于一条(已删除(注释，我已将代码更新为使用torch.reshape。

Answer 1

解决方案1作为一种可能的解决方案，您可以从输入x和x.shape[0]中获得批量大小，并将其用于以后的整形

import torch
batch = torch.zeros([128, 3, 96, 128], dtype=torch.float32)
# -1 will compute last dimension automatically
batch_upd = torch.reshape(batch, (batch.shape[0], -1))
print(batch_upd.shape)

此代码的输出为

torch.Size([128, 36864])

解决方案2作为另一种可能的解决方案，您可以使用压平

batch_upd = batch.flatten(start_dim=1)

将产生相同的输出

至于您的下一个问题，请考虑使用修改后的forward代码：

def forward(self, x):
x = x.flatten(1)  # shape: [B, C, W, H] -> [B, C*W*H]
x = F.relu(self.hidden(x))      # activation function for hidden layer
x = self.predict(x)             # linear output
x = x.reshape((-1, self.Chn, self.W, self.H)) # shape: [B, C*W*H] -> [B, C, W, H]
return x

以下是成功的使用示例：

ln = LinearNet(3, 96, 128, 256)
batch = torch.zeros((128, 3, 96, 128))
res = ln(batch)
print(res.shape)  # torch.Size([128, 3, 96, 128])