forked from dinhgit/FL-GAN_COVID
-
Notifications
You must be signed in to change notification settings - Fork 0
/
model_GAN3.py
123 lines (100 loc) · 3.75 KB
/
model_GAN3.py
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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
# -*- coding: utf-8 -*-
"""
Created on Sun Dec 13 19:37:59 2020
This is GAN model used for COVID data augmentation
@author: cdnguyen
"""
import torch.nn as nn
import torch.nn.functional as F
import torch
class netG(nn.Module):
def __init__(self, nz, ngf, nc):
super(netG, self).__init__()
self.ReLU = nn.ReLU(True)
self.Tanh = nn.Tanh()
#self.DropOut = nn.Dropout(p=0.75)
#self.conv0 = nn.ConvTranspose2d(nz, ngf * 16, 4, 1, 1, bias=False)
#self.BatchNorm0 = nn.BatchNorm2d(ngf * 16)
self.conv1 = nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False)
self.BatchNorm1 = nn.BatchNorm2d(ngf * 8)
self.conv2 = nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False)
self.BatchNorm2 = nn.BatchNorm2d(ngf * 4)
self.conv3 = nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False)
self.BatchNorm3 = nn.BatchNorm2d(ngf * 2)
self.conv4 = nn.ConvTranspose2d(ngf * 2, ngf * 1, 4, 2, 1, bias=False)
self.BatchNorm4 = nn.BatchNorm2d(ngf * 1)
self.conv5 = nn.ConvTranspose2d(ngf * 1, nc, 4, 2, 1, bias=False)
self.apply(weights_init)
def forward(self, input):
#x = self.conv0(input)
#x = self.BatchNorm0(x)
#x = self.ReLU(x)
x = self.conv1(input)
x = self.BatchNorm1(x)
x = self.ReLU(x)
#x = self.DropOut(x)
x = self.conv2(x)
x = self.BatchNorm2(x)
x = self.ReLU(x)
#x = self.DropOut(x)
x = self.conv3(x)
x = self.BatchNorm3(x)
x = self.ReLU(x)
#x = self.DropOut(x)
x = self.conv4(x)
x = self.BatchNorm4(x)
x = self.ReLU(x)
#x = self.DropOut(x)
x = self.conv5(x)
output = self.Tanh(x)
return output
class netD(nn.Module):
def __init__(self, ndf, nc, nb_label):
super(netD, self).__init__()
self.LeakyReLU = nn.LeakyReLU(0.2, inplace=True)
self.DropOut1 = nn.Dropout(p=0.5)
self.DropOut2 = nn.Dropout(p=0.25)
self.conv1 = nn.Conv2d(nc, ndf, 4, 2, 1, bias=False)
self.conv2 = nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False)
self.BatchNorm2 = nn.BatchNorm2d(ndf * 2)
self.conv3 = nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False)
self.BatchNorm3 = nn.BatchNorm2d(ndf * 4)
self.conv4 = nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False)
self.BatchNorm4 = nn.BatchNorm2d(ndf * 8)
self.conv5 = nn.Conv2d(ndf * 8, ndf * 1, 4, 1, 0, bias=False)
self.disc_linear = nn.Linear(ndf * 1, 1)
self.aux_linear = nn.Linear(ndf * 1, nb_label)
self.softmax = nn.Softmax()
self.sigmoid = nn.Sigmoid()
self.ndf = ndf
self.apply(weights_init)
def forward(self, input):
x = self.conv1(input)
x = self.LeakyReLU(x)
x = self.DropOut1(x)
x = self.conv2(x)
x = self.BatchNorm2(x)
x = self.LeakyReLU(x)
#x = self.DropOut(x)
x = self.conv3(x)
x = self.BatchNorm3(x)
x = self.LeakyReLU(x)
x = self.DropOut1(x)
x = self.conv4(x)
x = self.BatchNorm4(x)
x = self.LeakyReLU(x)
x = self.DropOut2(x)
x = self.conv5(x)
x = x.view(-1, self.ndf * 1)
c = self.aux_linear(x)
c = self.softmax(c)
s = self.disc_linear(x)
s = self.sigmoid(s)
return s,c
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)