Search code examples
pythonmachine-learningcomputer-visiongenerative-adversarial-network

Why and How to make range of Output image from a Neural Network [-1,1] in SRGAN?


I've been working on reimplementing Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network (SRGAN), now I'm stuck with the given information in section 3.2.enter image description here Based on the paper, the target HR should be in the range [-1,1], the input LR should be in the range [0,1], and MSE loss was calculated on images in range [-1,1].

The last sentence implies that the output from the generator network should be in [-1,1]. So that the content_loss = MSELoss(generator(output), target). Am I understanding correctly? But when I print the output from my generator network, whose laster is just a conv2d, it gives me images in the rage [0,1].

I'm not getting a good result by running SRResNet-MSE part, and I think maybe because the MSE loss is calculating on different ranges instead of just one [-1,1]?

But how can the output from my generator be in range [-1,1] if I still want to follow paper's architecture, which has conv2d as the last layer? enter image description here

I also include my code Class Generator here

class Generator(nn.Module):
    def __init__(self, scale_factor):
        upsample_block_num = int(math.log(scale_factor, 2))

        super(Generator, self).__init__()
        self.block1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=9, stride=1, padding=4, bias=False),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.block2 = self.make_layer(ResidualBlock, 16)
        self.block3 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.InstanceNorm2d(64, affine=True)
        )
        block4 = [UpsampleBLock(64, 2) for _ in range(upsample_block_num)]
        block4.append(nn.Conv2d(64, 3, kernel_size=9, stride=1, padding=4, bias=False))
        self.block4 = nn.Sequential(*block4)

    def make_layer(self, block, num_of_layer):
        layers = []
        for _ in range(num_of_layer):
            layers.append(block(64))
        return nn.Sequential(*layers)

    def forward(self, x):
        block1 = self.block1(x)
        block2 = self.block2(block1)
        block3 = self.block3(block2)
        block4 = self.block4(block1 + block3)

        return block4

Thank you so much for helping!


Solution

  • Given that you want output in the range [-1,1] and currently you are getting output in [0,1] simply doing.

    output = (output*2)-1
    

    should do the job.