如何将pytorch模型输出归一化在[0,1]范围内

Question

假设我有一个名为 UNet 的模型

output = UNet(input)

该输出是灰度图像形状的向量：(batch_size,1,128,128)

我想要做的是将每个图像标准化为[0,1] 范围内。

我是这样做的：

for i in range(batch_size):
   output[i,:,:,:] = output[i,:,:,:]/torch.amax(output,dim=(1,2,3))[i]

现在输出中的每个图像都已标准化，但是当我训练这样的模型时，pytorch 声称它无法计算此过程中的梯度，我明白为什么。

我的问题是在不破坏反向传播流的情况下标准化图像的正确方法是什么？ 像

output = UNet(input)
output = output.normalize
output2 = some_model(output)
loss = ..
loss.backward()
optimize.step()

我现在唯一的选择是在 UNet 末尾添加一个 sigmoid 激活，但我认为这不是一个好主意..

更新 - 代码（gen2,disc = unet,discriminator models.est_bias 是一些输出）：

---

更新 2x 代码：

with torch.no_grad():
            est_bias_for_disc = gen2(input_img)

            est_bias_for_disc /= est_bias_for_disc.amax(dim=(1,2,3), keepdim=True)
        disc_fake_hat = disc(est_bias_for_disc.detach())
        disc_fake_loss = BCE(disc_fake_hat, torch.zeros_like(disc_fake_hat)) 

        disc_real_hat = disc(bias_ref)
        disc_real_loss = BCE(disc_real_hat, torch.ones_like(disc_real_hat))
        disc_loss = (disc_fake_loss + disc_real_loss) / 2
        if epoch<=epochs_till_gen2_stop:
            disc_loss.backward(retain_graph=True) # Update gradients    
            opt_disc.step() # Update optimizer

然后有单独的培训：

 opt_gen2.zero_grad()
 est_bias = gen2(input_img)
 est_bias /= est_bias.amax(dim=(1,2,3), keepdim=True)
 disc_fake = disc(est_bias) 
 ADV_loss = BCE(disc_fake, torch.ones_like(disc_fake))
 gen2_loss = ADV_loss
 gen2_loss.backward() 
 opt_gen2.step()

Answer 1

由于批量维度上的索引，您正在覆盖张量的值。相反，您可以以矢量化形式执行操作：

output = output / output.amax(dim=(1,2,3), keepdim=True)

keepdim=True 参数保持torch.Tensor.amax 的输出形状与其输入的形状相同，允许您使用它执行就地操作。

Answer 2

您可以使用归一化功能：

>>> import torch
>>> import torch.nn.functional as F
>>> x = torch.tensor([[3.,4.],[5.,6.],[7.,8.]])
>>> x = F.normalize(x, dim = 0)
>>> print(x)

  tensor([[0.3293, 0.3714],
          [0.5488, 0.5571],
          [0.7683, 0.7428]])

只要不使用out，这将给出一个可微的张量。