-
Notifications
You must be signed in to change notification settings - Fork 38
/
Model_alphaGAN.py
106 lines (83 loc) · 3.78 KB
/
Model_alphaGAN.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
import numpy as np
import torch
import os
from torch import nn
from torch import optim
from torch.nn import functional as F
#***********************************************
#Encoder and Discriminator has same architecture
#***********************************************
class Discriminator(nn.Module):
def __init__(self, channel=512,out_class=1,is_dis =True):
super(Discriminator, self).__init__()
self.is_dis=is_dis
self.channel = channel
n_class = out_class
self.conv1 = nn.Conv3d(1, channel//8, kernel_size=4, stride=2, padding=1)
self.conv2 = nn.Conv3d(channel//8, channel//4, kernel_size=4, stride=2, padding=1)
self.bn2 = nn.BatchNorm3d(channel//4)
self.conv3 = nn.Conv3d(channel//4, channel//2, kernel_size=4, stride=2, padding=1)
self.bn3 = nn.BatchNorm3d(channel//2)
self.conv4 = nn.Conv3d(channel//2, channel, kernel_size=4, stride=2, padding=1)
self.bn4 = nn.BatchNorm3d(channel)
self.conv5 = nn.Conv3d(channel, n_class, kernel_size=4, stride=1, padding=0)
def forward(self, x, _return_activations=False):
h1 = F.leaky_relu(self.conv1(x), negative_slope=0.2)
h2 = F.leaky_relu(self.bn2(self.conv2(h1)), negative_slope=0.2)
h3 = F.leaky_relu(self.bn3(self.conv3(h2)), negative_slope=0.2)
h4 = F.leaky_relu(self.bn4(self.conv4(h3)), negative_slope=0.2)
h5 = self.conv5(h4)
if self.is_dis:
output = F.sigmoid(h5.view(h5.size()[0],-1))
else:
output = h5.view(h5.size()[0],-1)
return output
class Code_Discriminator(nn.Module):
def __init__(self, code_size=100,num_units=750):
super(Code_Discriminator, self).__init__()
n_class = 1
self.l1 = nn.Sequential(nn.Linear(code_size, num_units),
nn.BatchNorm1d(num_units),
nn.LeakyReLU(0.2,inplace=True))
self.l2 = nn.Sequential(nn.Linear(num_units, num_units),
nn.BatchNorm1d(num_units),
nn.LeakyReLU(0.2,inplace=True))
self.l3 = nn.Linear(num_units, 1)
def forward(self, x):
h1 = self.l1(x)
h2 = self.l2(h1)
h3 = self.l3(h2)
output = F.sigmoid(h3)
return output
class Generator(nn.Module):
def __init__(self, noise:int=100, channel:int=64):
super(Generator, self).__init__()
_c = channel
self.relu = nn.ReLU()
self.noise = noise
self.tp_conv1 = nn.ConvTranspose3d(noise, _c*8, kernel_size=4, stride=1, padding=0, bias=False)
self.bn1 = nn.BatchNorm3d(_c*8)
self.tp_conv2 = nn.Conv3d(_c*8, _c*4, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm3d(_c*4)
self.tp_conv3 = nn.Conv3d(_c*4, _c*2, kernel_size=3, stride=1, padding=1, bias=False)
self.bn3 = nn.BatchNorm3d(_c*2)
self.tp_conv4 = nn.Conv3d(_c*2, _c, kernel_size=3, stride=1, padding=1, bias=False)
self.bn4 = nn.BatchNorm3d(_c)
self.tp_conv5 = nn.Conv3d(_c, 1, kernel_size=3, stride=1, padding=1, bias=False)
def forward(self, noise):
noise = noise.view(-1,self.noise,1,1,1)
h = self.tp_conv1(noise)
h = self.relu(self.bn1(h))
h = F.upsample(h,scale_factor = 2)
h = self.tp_conv2(h)
h = self.relu(self.bn2(h))
h = F.upsample(h,scale_factor = 2)
h = self.tp_conv3(h)
h = self.relu(self.bn3(h))
h = F.upsample(h,scale_factor = 2)
h = self.tp_conv4(h)
h = self.relu(self.bn4(h))
h = F.upsample(h,scale_factor = 2)
h = self.tp_conv5(h)
h = F.tanh(h)
return h