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

How to normalize pytorch model output to be in range [0,1]

假设我有一个名为 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,鉴别器模型。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()

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

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

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

您可以使用归一化函数:

>>> 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,这将给出一个可微张量。