-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
94 lines (77 loc) · 5.05 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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
import keras.backend as K
import tensorflow as tf
from keras.layers import Input, Conv2D, BatchNormalization, UpSampling2D
from keras.models import Model
from keras.regularizers import l2
from keras.utils import multi_gpu_model
from keras.utils import plot_model
from config import img_rows, img_cols, num_classes, kernel
l2_reg = l2(1e-3)
def build_model():
input_tensor = Input(shape=(img_rows, img_cols, 1))
x = Conv2D(64, (kernel, kernel), activation='relu', padding='same', name='conv1_1', kernel_initializer="he_normal",
kernel_regularizer=l2_reg)(input_tensor)
x = Conv2D(64, (kernel, kernel), activation='relu', padding='same', name='conv1_2', kernel_initializer="he_normal",
kernel_regularizer=l2_reg, strides=(2, 2))(x)
x = BatchNormalization()(x)
x = Conv2D(128, (kernel, kernel), activation='relu', padding='same', name='conv2_1', kernel_initializer="he_normal",
kernel_regularizer=l2_reg)(x)
x = Conv2D(128, (kernel, kernel), activation='relu', padding='same', name='conv2_2', kernel_initializer="he_normal",
kernel_regularizer=l2_reg,
strides=(2, 2))(x)
x = BatchNormalization()(x)
x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv3_1',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv3_2',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv3_3', kernel_initializer="he_normal",
strides=(2, 2))(x)
x = BatchNormalization()(x)
x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', name='conv4_1',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', name='conv4_2',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', name='conv4_3',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = BatchNormalization()(x)
x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv5_1',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv5_2',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv5_3',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = BatchNormalization()(x)
x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv6_1',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv6_2',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv6_3',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = BatchNormalization()(x)
x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv7_1',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv7_2',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv7_3',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = BatchNormalization()(x)
x = UpSampling2D(size=(2, 2))(x)
x = Conv2D(128, (kernel, kernel), activation='relu', padding='same', name='conv8_1',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(128, (kernel, kernel), activation='relu', padding='same', name='conv8_2',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = Conv2D(128, (kernel, kernel), activation='relu', padding='same', name='conv8_3',
kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
x = BatchNormalization()(x)
outputs = Conv2D(num_classes, (1, 1), activation='softmax', padding='same', name='pred')(x)
model = Model(inputs=input_tensor, outputs=outputs, name="ColorNet")
return model
if __name__ == '__main__':
with tf.device("/cpu:0"):
encoder_decoder = build_model()
print(encoder_decoder.summary())
plot_model(encoder_decoder, to_file='encoder_decoder.svg', show_layer_names=True, show_shapes=True)
parallel_model = multi_gpu_model(encoder_decoder, gpus=None)
print(parallel_model.summary())
plot_model(parallel_model, to_file='parallel_model.svg', show_layer_names=True, show_shapes=True)
K.clear_session()