taesungp / contrastive-unpaired-translation

Contrastive unpaired image-to-image translation, faster and lighter training than cyclegan (ECCV 2020, in PyTorch)
https://taesung.me/ContrastiveUnpairedTranslation/
Other
2.23k stars 417 forks source link

Change upsample method in ResnetGenerator class #157

Open bekhzod-olimov opened 1 year ago

bekhzod-olimov commented 1 year ago

When I tried to change the Upsample method in the ResnetGenerator class to Upsample2, it suddenly requires more than 500GB and raises an error. What is the reason for this?

Here is the code (in models/networks.py):

else:
      model += [Upsample2(ngf * mult),
                nn.Conv2d(ngf * mult, int(ngf * mult / 2),
                          kernel_size=3, stride=1,
                          padding=1,  # output_padding=1,
                          bias=use_bias),
                norm_layer(int(ngf * mult / 2)),
                nn.ReLU(True)]

And this is the error:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00 GiB (GPU 0; 39.59 GiB total capacity; 885.03 MiB already allocated; 20.78 GiB free; 1.89 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

bekhzod-olimov commented 1 year ago

Oh, I guess the reason was the argument to the Upsamle2 function (ngf * mult) is a very big number so it raised an error. Changing it to scale_factor of 2 resolved the issue. But I have one more question, by default ResnetGenerator uses Conv2dTranspose

class Upsample(nn.Module):
    def __init__(self, channels, pad_type='repl', filt_size=4, stride=2):
        super(Upsample, self).__init__()
        self.filt_size = filt_size
        self.filt_odd = np.mod(filt_size, 2) == 1
        self.pad_size = int((filt_size - 1) / 2)
        self.stride = stride
        self.off = int((self.stride - 1) / 2.)
        self.channels = channels

        filt = get_filter(filt_size=self.filt_size) * (stride**2)
        self.register_buffer('filt', filt[None, None, :, :].repeat((self.channels, 1, 1, 1)))

        self.pad = get_pad_layer(pad_type)([1, 1, 1, 1])

    def forward(self, inp):
        ret_val = F.conv_transpose2d(self.pad(inp), self.filt, stride=self.stride, padding=1 + self.pad_size, groups=inp.shape[1])[:, :, 1:, 1:]
        if(self.filt_odd):
            return ret_val
        else:
            return ret_val[:, :, :-1, :-1]

and changing it to Upsample2d method must decrease the number of trainable parameters, right?