SpacEX/generator.py at main · ahmedemam576/SpacEX · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
 #!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Aug  5 16:09:09 2022

@author: ahmedemam576
"""
import torch
from torch import nn

from expanding_block import Expanding_Block
#the expanding block divides the number of the filters by two, output_filters = input_filters//2
from contracting_block import Contracting_Block
#the contracting block doubles the number of the filters, output_filters = 2 * input_filters

from residual_block import Residual_Block
# the residual block consists of 2 convolution layers, and at the end we concatenate the input to the output
# also the input filters and the output filters are equal
from feature_map import Feature_map
# the feature map class doesn't change the number of the filters , it's just a convolutional layer

class Generator(nn.Module):
    def __init__(self, input_channels, output_channels, hidden_channels=64):

        super(Generator, self).__init__()
        self.up_feature =Feature_map(input_channels, hidden_channels)


        self.contract1 = Contracting_Block(hidden_channels)

        self.contract2 = Contracting_Block(hidden_channels*2)

        self.res1 = Residual_Block(hidden_channels*4)
        self.res2 = Residual_Block(hidden_channels*4)
        self.res3 = Residual_Block(hidden_channels*4)
        self.res4 = Residual_Block(hidden_channels*4)
        self.res5 = Residual_Block(hidden_channels*4)
        self.res6 = Residual_Block(hidden_channels*4)
        self.res7 = Residual_Block(hidden_channels*4)
        self.res8 = Residual_Block(hidden_channels*4)
        self.res9 = Residual_Block(hidden_channels*4)

        self.expand1 = Expanding_Block(hidden_channels*4)
        self.expand2 = Expanding_Block(hidden_channels*2)

        self.down_feature = Feature_map(hidden_channels, output_channels)
        self.tanh = nn.Tanh()
        self.init_weights()


    # initialize the parmaters weights
    def init_weights(self):

        for m in self.modules():

            #print(m)
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                print('yes conv2d')
                torch.nn.init.xavier_normal_(m.weight)
            if isinstance(m, nn.BatchNorm2d):
                torch.nn.init.normal_(m.weight, 0.0, 0.0)
                torch.nn.init.constant_(m.bias, 0)
            if isinstance(m, nn.BatchNorm2d):
                print('yes batch norm')
                # these is a problem here
                print(m)
                torch.nn.init.normal_(m.weight, 0.0, 0.0)
                torch.nn.init.constant_(m.bias, 0)


    def forward(self, x):
        x= self.up_feature(x)
        x= self.contract1(x)
        x= self.contract2(x)
        x= self.res1(x)
        x= self.res2(x)
        x= self.res3(x)
        x= self.res4(x)
        x= self.res5(x)
        x= self.res6(x)
        x= self.res7(x)
        x= self.res8(x)
        x= self.res9(x)
        x= self.expand1(x)
        x= self.expand2(x)
        x= self.down_feature(x)
        x= self.tanh(x)
        return x