-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
73 lines (60 loc) · 3.08 KB
/
model.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
'''
Copyright (C) 2020, Northwestern University and Lawrence Berkeley National Laboratory
See COPYRIGHT notice in top-level directory.
'''
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input
from tensorflow.keras.layers import Conv3D
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import AveragePooling3D
from tensorflow.keras.layers import ReLU
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Dropout
from tensorflow.keras.layers import BatchNormalization
class model ():
def __init__ (self, num_filters = 16, kernel_size = 3, pool_size = 2):
self.num_filters = num_filters
self.kernel_size = kernel_size
self.pool_size = pool_size
print ("Initializing cosmoflow model...")
def build_model(self):
x_in = Input(shape = (128, 128, 128, 12), name = "input")
# The first batch normalization should be replaced with preprocessing later.
h = BatchNormalization()(x_in)
h = Conv3D(16, self.kernel_size, padding = 'same', activation = 'linear', name = 'conv1')(h)
h = BatchNormalization()(h)
h = ReLU(name = 'relu1')(h)
h = AveragePooling3D(pool_size = self.pool_size, name = 'pool1')(h)
h = Conv3D(32, self.kernel_size, padding = 'same', activation = 'linear', name = 'conv2')(h)
h = BatchNormalization()(h)
h = ReLU(name = 'relu2')(h)
h = AveragePooling3D(pool_size = self.pool_size, name = 'pool2')(h)
h = Conv3D(64, self.kernel_size, padding = 'same', activation = 'linear', name = 'conv3')(h)
h = BatchNormalization()(h)
h = ReLU(name = 'relu3')(h)
h = AveragePooling3D(pool_size = self.pool_size, name = 'pool3')(h)
h = Conv3D(128, self.kernel_size, strides = (2, 2, 2), padding = 'same', activation = 'linear', name = 'conv4')(h)
h = BatchNormalization()(h)
h = ReLU(name = 'relu4')(h)
h = AveragePooling3D(pool_size = self.pool_size, name = 'pool4')(h)
h = Conv3D(256, self.kernel_size, padding = 'same', activation = 'linear', name = 'conv5')(h)
h = BatchNormalization()(h)
h = ReLU(name = 'relu5')(h)
h = AveragePooling3D(pool_size = self.pool_size, name = 'pool5')(h)
h = Conv3D(256, self.kernel_size, padding = 'same', activation = 'linear', name = 'conv6')(h)
h = BatchNormalization()(h)
h = ReLU(name = 'relu6')(h)
h = Conv3D(256, self.kernel_size, padding = 'same', activation = 'linear', name = 'conv7')(h)
h = BatchNormalization()(h)
h = ReLU(name = 'relu7')(h)
h = Flatten(name = 'pool_to_full')(h)
h = Dropout(0.05, name='drop1')(h)
h = Dense(2048, activation = 'linear', name = 'full1')(h)
h = Dropout(0.05, name='drop2')(h)
h = ReLU(name = 'relu8')(h)
h = Dense(256, activation = 'linear', name = 'full2')(h)
h = Dropout(0.05, name='drop3')(h)
h = ReLU(name = 'relu9')(h)
y = Dense(4, activation = 'linear', name = 'full3')(h)
return Model(inputs = x_in, outputs = y)