155 lines
3.8 KiB
Python
155 lines
3.8 KiB
Python
import tensorflow as tf
|
|
import numpy as np
|
|
|
|
from tensorflow import keras as k
|
|
from tensorflow.keras import layers as kl
|
|
from tensorflow.keras import models as km
|
|
from tensorflow.keras import optimizers as ko
|
|
|
|
CONV_SIZE = 128
|
|
MPPT_SIZE = 1024
|
|
DRPT_RATE = 0.3
|
|
TOUT_AMNT = 2
|
|
|
|
def conv(i, n, s, t = (1, 1)):
|
|
c = kl.Conv2D(n, s, padding = 'same', strides = t)(i)
|
|
b = kl.BatchNormalization()(c)
|
|
return kl.Activation('relu')(b)
|
|
|
|
def avg(i, s):
|
|
return kl.AveragePooling2D(s, padding = 'same', strides = (1, 1))(i)
|
|
|
|
def max(i, s):
|
|
return kl.MaxPooling2D(s, padding = 'same', strides = (1, 1))(i)
|
|
|
|
def generate_start(i):
|
|
r1 = conv(i, CONV_SIZE, (3, 3))
|
|
r2 = conv(r1, CONV_SIZE, (3, 3))
|
|
r3 = conv(r2, CONV_SIZE, (3, 3))
|
|
|
|
m1 = kl.MaxPooling2D((3, 3))(r3)
|
|
|
|
r4 = conv(m1, CONV_SIZE, (1, 1))
|
|
r5 = conv(r4, CONV_SIZE, (3, 3))
|
|
|
|
return kl.MaxPooling2D((3, 3))(r5)
|
|
|
|
def generate_type_a(i):
|
|
a11 = conv(i, CONV_SIZE, (1, 1), (3, 3))
|
|
|
|
a21 = avg(i, (3, 3))
|
|
a22 = conv(a21, CONV_SIZE, (1, 1))
|
|
|
|
a31 = conv(i, CONV_SIZE, (1, 1))
|
|
a32 = conv(a31, CONV_SIZE, (3, 3), (3, 3))
|
|
|
|
a41 = conv(i, CONV_SIZE, (1, 1))
|
|
a42 = conv(a41, CONV_SIZE, (3, 3))
|
|
a43 = conv(a42, CONV_SIZE, (3, 3), (3, 3))
|
|
|
|
return kl.Concatenate()([a11, a22, a32, a43])
|
|
|
|
def generate_ab_bridge(i):
|
|
ab11 = conv(i, CONV_SIZE, (1, 1))
|
|
ab12 = conv(ab11, CONV_SIZE, (3, 3))
|
|
ab13 = conv(ab12, CONV_SIZE, (3, 3), (3, 3))
|
|
|
|
ab21 = conv(i, CONV_SIZE, (3, 3), (3, 3))
|
|
|
|
ab31 = max(i, (3, 3))
|
|
|
|
return kl.Concatenate()([ab13, ab21, ab31])
|
|
|
|
def generate_type_b(i):
|
|
b11 = conv(i, CONV_SIZE, (1, 1))
|
|
b12 = conv(b11, CONV_SIZE, (7, 1))
|
|
b13 = conv(b12, CONV_SIZE, (1, 7))
|
|
b14 = conv(b13, CONV_SIZE, (7, 1))
|
|
b15 = conv(b14, CONV_SIZE, (1, 7), (3, 3))
|
|
|
|
b21 = conv(i, CONV_SIZE, (1, 1))
|
|
b22 = conv(b21, CONV_SIZE, (1, 7))
|
|
b23 = conv(b22, CONV_SIZE, (7, 1), (3, 3))
|
|
|
|
b31 = conv(i, CONV_SIZE, (1, 1), (3, 3))
|
|
|
|
b41 = avg(i, (3, 3))
|
|
b42 = conv(b41, CONV_SIZE, (1, 1))
|
|
|
|
return kl.Concatenate()([b15, b23, b31, b42])
|
|
|
|
def generate_bc_bridge(i):
|
|
bc11 = conv(i, CONV_SIZE, (1, 1))
|
|
bc12 = conv(bc11, CONV_SIZE, (7, 1))
|
|
bc13 = conv(bc12, CONV_SIZE, (1, 7))
|
|
bc14 = conv(bc13, CONV_SIZE, (3, 3), (3, 3))
|
|
|
|
bc21 = conv(i, CONV_SIZE, (1, 1))
|
|
bc22 = conv(bc21, CONV_SIZE, (1, 1), (3, 3))
|
|
|
|
bc31 = max(i, (3, 3))
|
|
|
|
return kl.Concatenate()([bc14, bc22, bc31])
|
|
|
|
def generate_aux(i):
|
|
u1 = kl.AveragePooling2D((5, 5))(i)
|
|
u2 = conv(u1, CONV_SIZE, (1, 1))
|
|
u3 = kl.Flatten()(u2)
|
|
u4 = kl.Dropout(DRPT_RATE)(u3)
|
|
u5 = kl.Dense(MPPT_SIZE)(u4)
|
|
|
|
return kl.Dense(TOUT_AMNT)(u5)
|
|
|
|
def generate_type_c(i):
|
|
c11 = conv(i, CONV_SIZE, (1, 1))
|
|
c12 = conv(c11, CONV_SIZE, (3, 3))
|
|
c13 = conv(c12, CONV_SIZE, (1, 3), (3, 3))
|
|
c14 = conv(c12, CONV_SIZE, (3, 1), (3, 3))
|
|
|
|
c21 = conv(i, CONV_SIZE, (1, 1))
|
|
c22 = conv(c21, CONV_SIZE, (1, 3), (3, 3))
|
|
c23 = conv(c21, CONV_SIZE, (3, 1), (3, 3))
|
|
|
|
c31 = max(i, (3, 3))
|
|
c32 = conv(c31, CONV_SIZE, (1, 1))
|
|
|
|
c41 = conv(i, CONV_SIZE, (1, 1), (3, 3))
|
|
|
|
return kl.Concatenate()([c14, c23, c32, c41])
|
|
|
|
def generate_finish(i):
|
|
f1 = kl.AveragePooling2D((5, 5))(i)
|
|
f2 = conv(f1, CONV_SIZE, (1, 1))
|
|
f3 = kl.Flatten()(f2)
|
|
f4 = kl.Dropout(DRPT_RATE)(f3)
|
|
f5 = kl.Dense(MPPT_SIZE)(f4)
|
|
|
|
return kl.Dense(TOUT_AMNT)(f5)
|
|
|
|
|
|
gi = kl.Input((2048, 2048, 3))
|
|
|
|
ds = generate_start(gi)
|
|
|
|
for _ in range(1):
|
|
ds = generate_type_a(ds)
|
|
|
|
ds = generate_ab_bridge(ds)
|
|
|
|
for _ in range(1):
|
|
ds = generate_type_b(ds)
|
|
|
|
uo = generate_aux(ds)
|
|
|
|
go = generate_bc_bridge(ds)
|
|
|
|
for _ in range(1):
|
|
go = generate_type_c(go)
|
|
|
|
go = generate_finish(go)
|
|
|
|
mod = k.Model(inputs = gi, outputs = [go, uo])
|
|
|
|
mod.summary()
|
|
mod.compile(optimizer = ko.Lion(learning_rate = 0.001), metrics = ['accuracy', 'val_accuracy'])
|