如何将50×50图像红通道转为扩散模型FC网络的2维输入？

我为扩散模型构建了如下全连接网络NetD，此前使用形状为[3000,2]的均匀分布数据完成训练，现在希望改用尺寸为(50,50)的图像红通道数据训练该模型，但不确定如何将[50,50]形状的张量转换为符合网络输入要求的2维格式，特此咨询正确的维度转换方法。

去噪网络NetD代码 #

import torch
from torch import linalg as LA
import plotly.express as px
from plotly.subplots import make_subplots
import plotly.graph_objects as go
from torch.utils.data import TensorDataset, DataLoader, Dataset
import torch.optim as optim
import torch.nn as nn
import matplotlib.pyplot as plt
from tqdm import tqdm

lr = 0.0001
num_data = 3000
batch_size = 3000
num_epochs = 3000
x_dim = 2  # 输入维度
beta1 = 0.5  # Adam优化器的Beta1超参数
l2 = nn.MSELoss()
colors = px.colors.qualitative.T10
MODEL_PATH = 'denoiser.pth'
CONDITIONAL_MODEL_PATH = 'denoiser_conditional.pth'

# 去噪网络
class NetD(nn.Module):
    def __init__(self):
        super(NetD, self).__init__()
        # 输入: [x0, x1, t]
        self.fc = nn.Sequential(
            nn.Linear(x_dim + 1, (x_dim + 1) * 8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear((x_dim + 1) * 8, (x_dim + 1) * 20),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear((x_dim + 1) * 20, x_dim),
            # 输出: [e0, e1]
        )

    def forward(self, x_t):
        return self.fc(x_t)

图像红通道处理代码 #

我已经完成了图像红通道的提取与预处理，代码如下：

import torchvision.transforms as transforms
import torchvision.models as models
from PIL import Image

# 输出图像目标尺寸
imsize = (50,50)  # 无GPU时使用小尺寸

loader = transforms.Compose([
    transforms.Resize(imsize),  # 缩放图像
    transforms.ToTensor()])  # 转换为torch张量

def image_loader(image_name):
    image = Image.open(image_name)
    image = loader(image).unsqueeze(0)
    return image.to(device, torch.float)

cat = image_loader('/content/orange_cat.jpg')
red_channel = cat[:, 0, :, :]  # 提取图像的红通道（R）
density = torch.where(red_channel > 0.78, torch.tensor(1.0), torch.tensor(0.0))
data_1 = red_channel.squeeze().to(device)

当前尝试与疑问 #

我尝试用以下代码转换维度：

# 将张量重塑为2维张量
reshaped_red_channel = data_1.view(-1, 2)

转换后得到尺寸为[1250,2]的张量，虽然元素总数与原张量一致，但不确定这种转换方式是否符合扩散模型训练的要求，需要正确的维度转换方法。

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