junyanz / pytorch-CycleGAN-and-pix2pix

Image-to-Image Translation in PyTorch
Other
22.83k stars 6.29k forks source link

Switch from GPU to CPU why, perceptual loss CycleGAN training #822

Open Scienceseb opened 4 years ago

Scienceseb commented 4 years ago

Problem: model always switch from GPU to CPU

Modifications done for adding the perceptual loss: So if you want to use perceptual loss while using this amazing CycleGAN do the following: In models.cycle_gan_model.py: add at the beginning:

from torchvision import models
import torch.nn as nn

class VGGNet(nn.Module):
def __init__(self):
    """Select conv1_1 ~ conv5_1 activation maps."""
    super(VGGNet, self).__init__()
    self.select = ['9', '36']
    self.vgg = models.vgg19(pretrained=True).features

def forward(self, x):
    """Extract multiple convolutional feature maps."""
    features = []
    for name, layer in self.vgg._modules.items():
        x = layer(x)
        if name in self.select:
            features.append(x)
    return features[0], features[1]

and in backward_G(self):

perceptual=True #False if you dont want perceptual loss in that training

 if perceptual==True:

        vgg = VGGNet().cuda().eval()
        with torch.no_grad():

            A=self.real_A
            B=self.real_B

            c = nn.MSELoss()

            rx = self.netG_B(self.netG_A(A))
            ry = self.netG_A(self.netG_B(B))

            fx1, fx2 = vgg(A)
            fy1, fy2 = vgg(B)

            frx1, frx2 = vgg(rx)
            fry1, fry2 = vgg(ry)

            m1 = c(fx1, frx1)
            m2 = c(fx2, frx2)

            m3 = c(fy1, fry1)
            m4 = c(fy2, fry2)

            self.perceptual_loss = (m1 + m2 + m3 + m4) * 0.00001 * 0.5
    else:
        self.perceptual_loss=0

and finally:

self.loss_G = self.loss_G_A + self.loss_G_B + self.loss_cycle_A + self.loss_cycle_B + self.loss_idt_A + self.loss_idt_B+self.perceptual_loss

You now have a perceptual cycle gan, congratulation!

But with those modifications the model always switch from GPU to CPU why?

junyanz commented 4 years ago

Here is how we initialize the network with GPU ids.

axelaco commented 4 years ago

@Scienceseb Hi, can you explain the role of this perceptual loss in your training ?, you use it to improve the generators ?

zzhan127 commented 4 years ago

It still uses GPU for training. However, since you call vgg = VGGNet().cuda().eval() in the backward function, it takes a long time for initialization each time. Try put this in the initialization function.

Scienceseb commented 4 years ago

@Scienceseb Hi, can you explain the role of this perceptual loss in your training ?, you use it to improve the generators ?

Perceptual loss make the resulting images sharper.

Sun123-fei commented 3 years ago

@Scienceseb Hi, can you explain the role of this perceptual loss in your training ?, you use it to improve the generators ?

Perceptual loss make the resulting images sharper.

Excuse me,what is the effect of adding perceptual loss

Sun123-fei commented 3 years ago

@Scienceseb Hi, can you explain the role of this perceptual loss in your training ?, you use it to improve the generators ?

Perceptual loss make the resulting images sharper.

how is it?

Youqi970124 commented 3 years ago

If I want save this perceptual loss to the log_loss.txt,and visualization it, where I need to change.

very thanks