Adding new layers between intermediate layers of a GAN

vision
1

I have a Generator defined as:

class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d( ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # state size. (ngf) x 32 x 32
            nn.ConvTranspose2d( ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )

    def forward(self, input):
        return self.main(input)

Between every ConvTranspose2d layer and BatchNorm2d layer I want to insert some kind of mathematical operation that is applied on the output of the conv layer.

For example, let’s say the output of the conv layer is X and i want to apply a shift b to this output such that X + b is now the input to the BN layer.

How can I do this?

2

You could create a new custom nn.Module and add it into the nn.Sequential container:

class ShiftLayer(nn.Module):
    def __init__(self, b):
        super().__init__()
        self.register_buffer('b', b)
        
    def forward(self, x):
        return x + self.b

class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False),
            ShiftLayer(b=torch.tensor([2.])),
            nn.BatchNorm2d(ngf * 8),
            ...

Registering b as a buffer will make sure that the to() operation such as model.to('cuda') will move this buffer to the specified device etc.
Also, it will be stored in the state_dict of the module so that you could save and load it properly.

3

Hi, that solution works, thanks! I’m now trying to modify the x and b parameters from this ShiftLayer based on some conditions

netG.x and netG.b does not work, it gives me the error

torch.nn.modules.module.ModuleAttributeError: 'Generator' object has no attribute 'x'

So I tried accessing them through netG.module.x but that does not work either and I get the error:

torch.nn.modules.module.ModuleAttributeError: 'Generator' object has no attribute 'module'

How can I do this?

4

ShiftLayer doesn’t have an x parameter, but you could access the b buffer via:

netG.main[1].b

assuming you are using the model definition posted in my code snippet.

5

Yes, sorry that was my bad, I am wanting to update b. So I actually tried what you suggested:

netG.main[1].b  += 5

and now the error is

'ShiftLayer' object has no attribute 'b'

Although I have initialized the generator exactly as you posted in your code snippet.

6

Okay, I looked over my code and it was a naming error on my end. It’s fixed now! thanks! :slight_smile:

7

Is it possible to define ShiftLayer so that I can train its parameter b with a torch.optim() function?

8

Yes, register b as an nn.Parameter and it’ll be trainable:

self.b = nn.Parameter(torch.tensor(...))
9

This has to be done when I define the ShiftLayer right?

class ShiftLayer(nn.Module):
    def __init__(self, b):
        super().__init__()
        self.register_buffer('b', b)
        self.b = nn.Parameter(torch.tensor(...))
        
    def forward(self, x):
        return x + self.b

like this?

10

Yes, you can define the parameter in the ShiftLayer, but would need to remove the self.register_buffer as it’s not needed.

11

So my final definition for ShiftLayer is

class ShiftLayer(nn.Module):
    def __init__(self, b):
        super().__init__()
        self.b = nn.Parameter(torch.tensor(...))
        
    def forward(self, x):
        return x + self.b

However, this gives me an error:

 self.b = nn.Parameter(torch.tensor(...)),                                                                                                ^                                                                                                                               
SyntaxError: expression cannot contain assignment, perhaps you meant "=="?

Not sure how else I would assign b as a parameter

12

The ... was just a placeholder for your actual values, so you would need to fill torch.tensor either with your values or use a factory method given a shape, e.g. torch.randn(10, 10).

13

yeah of couse. I have actual values like so -

self.b = nn.Parameter(torch.randn(32)),

still getting this error though :slight_smile:

14

I cannot reproduce the issue using:

class ShiftLayer(nn.Module):
    def __init__(self, b):
        super().__init__()
        self.b = nn.Parameter(torch.tensor(1.))
        
    def forward(self, x):
        return x + self.b
    
layer = ShiftLayer(1)
out = layer(torch.randn(1))

so could you post an executable code snippet raising this issue?

15

Sorry for the late reply, I fixed it! I had an accidental Sequential() command right after def __init__() which was the problem