PyTorch中UNet2DOutput无size属性报错的解决方法求助

错误原因 #

你使用的UNet2DModel（一般来自diffusers库）的前向传播返回的是UNet2DOutput类实例，而非直接的张量。这个实例里的sample属性才是模型输出的预测张量，而nn.MSELoss()需要两个张量来计算损失，直接传入UNet2DOutput对象就会触发AttributeError: 'UNet2DOutput' object has no attribute 'size'错误。

解决方法 #

修改模型输出的赋值语句，提取UNet2DOutput中的sample属性作为预测张量：

将：

pred = net(noisy_x, 0)

改为：

pred = net(noisy_x, 0).sample

修改后的完整代码 #

net = UNet2DModel(sample_size=28,
                    in_channels=1,
                    out_channels=1,
                    layers_per_block=2,
                    block_out_channels=(32, 64, 64),
                    down_block_types=(
                        "DownBlock2D",
                        "AttnDownBlock2D",
                        "AttnDownBlock2D",
                    ),
                    up_block_types=(
                        "AttnUpBlock2D",
                        "AttnUpBlock2D",
                        "UpBlock2D",
                    ),
)

batch_size = 64
n_epochs = 3
train_dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True)

net.to(device)

loss_fn = nn.MSELoss()

opt = torch.optim.Adam(net.parameters(), lr=1e-3)

losses = [] 

for epoch in range(n_epochs):
    for x, y in train_dataloader:
        x = x.to(device)
        noise_amount = torch.rand(x.shape[0]).to(device)
        
        noisy_x = corrupt(x, noise_amount)
        
        pred = net(noisy_x, 0).sample  # 此处为修改点
        
        loss = loss_fn(pred, x)
        
        opt.zero_grad()
        loss.backward()
        opt.step()
        
        losses.append(loss.item())

内容的提问来源于stack exchange，提问作者张佳伟