307 lines
14 KiB
Python
307 lines
14 KiB
Python
# taken from https://github.com/calmisential/TensorFlow2.0_ResNet/blob/master/models/resnet.py
|
|
|
|
import tensorflow as tf
|
|
|
|
|
|
class BasicBlock(tf.keras.layers.Layer):
|
|
|
|
def __init__(self, filter_num, stride=1, upsample=False):
|
|
super(BasicBlock, self).__init__()
|
|
if not upsample:
|
|
self.conv1 = tf.keras.layers.Conv2D(filters=filter_num,
|
|
kernel_size=(3, 3),
|
|
strides=stride,
|
|
padding="same")
|
|
else:
|
|
self.conv1 = tf.keras.layers.Conv2DTranspose(filters=filter_num,
|
|
kernel_size=(3, 3),
|
|
strides=stride,
|
|
padding="same")
|
|
self.bn1 = tf.keras.layers.BatchNormalization()
|
|
self.conv2 = tf.keras.layers.Conv2D(filters=filter_num,
|
|
kernel_size=(3, 3),
|
|
strides=1,
|
|
padding="same")
|
|
self.bn2 = tf.keras.layers.BatchNormalization()
|
|
if stride != 1:
|
|
if not upsample:
|
|
self.updownsample = tf.keras.Sequential()
|
|
self.updownsample.add(tf.keras.layers.Conv2D(filters=filter_num,
|
|
kernel_size=(1, 1),
|
|
strides=stride))
|
|
self.updownsample.add(tf.keras.layers.BatchNormalization())
|
|
else:
|
|
self.updownsample = tf.keras.Sequential()
|
|
self.updownsample.add(tf.keras.layers.Conv2DTranspose(filters=filter_num,
|
|
kernel_size=(1, 1),
|
|
strides=stride))
|
|
self.updownsample.add(tf.keras.layers.BatchNormalization())
|
|
else:
|
|
self.updownsample = lambda x: x
|
|
|
|
def call(self, inputs, training=None, **kwargs):
|
|
residual = self.updownsample(inputs)
|
|
|
|
x = self.conv1(inputs)
|
|
x = self.bn1(x, training=training)
|
|
x = tf.nn.relu(x)
|
|
x = self.conv2(x)
|
|
x = self.bn2(x, training=training)
|
|
|
|
output = tf.nn.relu(tf.keras.layers.add([residual, x]))
|
|
|
|
return output
|
|
|
|
|
|
class BottleNeck(tf.keras.layers.Layer):
|
|
def __init__(self, filter_num, stride=1, upsample=False):
|
|
super(BottleNeck, self).__init__()
|
|
self.conv1 = tf.keras.layers.Conv2D(filters=filter_num,
|
|
kernel_size=(1, 1),
|
|
strides=1,
|
|
padding='same')
|
|
self.bn1 = tf.keras.layers.BatchNormalization()
|
|
if not upsample:
|
|
self.conv2 = tf.keras.layers.Conv2D(filters=filter_num,
|
|
kernel_size=(3, 3),
|
|
strides=stride,
|
|
padding='same')
|
|
else:
|
|
self.conv2 = tf.keras.layers.Conv2DTranspose(filters=filter_num,
|
|
kernel_size=(3, 3),
|
|
strides=stride,
|
|
padding='same')
|
|
self.bn2 = tf.keras.layers.BatchNormalization()
|
|
self.conv3 = tf.keras.layers.Conv2D(filters=filter_num * 4,
|
|
kernel_size=(1, 1),
|
|
strides=1,
|
|
padding='same')
|
|
self.bn3 = tf.keras.layers.BatchNormalization()
|
|
|
|
self.updownsample = tf.keras.Sequential()
|
|
if not upsample:
|
|
self.updownsample.add(tf.keras.layers.Conv2D(filters=filter_num * 4,
|
|
kernel_size=(1, 1),
|
|
strides=stride))
|
|
else:
|
|
self.updownsample.add(tf.keras.layers.Conv2DTranspose(filters=filter_num * 4,
|
|
kernel_size=(1, 1),
|
|
strides=stride))
|
|
self.updownsample.add(tf.keras.layers.BatchNormalization())
|
|
|
|
def call(self, inputs, training=None, **kwargs):
|
|
residual = self.updownsample(inputs)
|
|
|
|
x = self.conv1(inputs)
|
|
x = self.bn1(x, training=training)
|
|
x = tf.nn.relu(x)
|
|
x = self.conv2(x)
|
|
x = self.bn2(x, training=training)
|
|
x = tf.nn.relu(x)
|
|
x = self.conv3(x)
|
|
x = self.bn3(x, training=training)
|
|
|
|
output = tf.nn.relu(tf.keras.layers.add([residual, x]))
|
|
|
|
return output
|
|
|
|
|
|
def make_basic_block_layer(filter_num, blocks, stride=1, upsample=False):
|
|
res_block = tf.keras.Sequential()
|
|
res_block.add(BasicBlock(filter_num, stride=stride, upsample=upsample))
|
|
|
|
for _ in range(1, blocks):
|
|
res_block.add(BasicBlock(filter_num, stride=1))
|
|
|
|
return res_block
|
|
|
|
|
|
def make_bottleneck_layer(filter_num, blocks, stride=1, upsample=False):
|
|
res_block = tf.keras.Sequential()
|
|
res_block.add(BottleNeck(filter_num, stride=stride, upsample=upsample))
|
|
|
|
for _ in range(1, blocks):
|
|
res_block.add(BottleNeck(filter_num, stride=1))
|
|
|
|
return res_block
|
|
|
|
|
|
class ResNetTypeI(tf.keras.Model):
|
|
def __init__(self, layer_params, features, reversed=False, labels=None, is_cvae=False):
|
|
super(ResNetTypeI, self).__init__()
|
|
self.reversed = reversed
|
|
self.labels = labels
|
|
self.layer_params = layer_params
|
|
|
|
if not reversed:
|
|
self.conv1 = tf.keras.layers.Conv2D(filters=32 if len(layer_params) > 1 else features,
|
|
kernel_size=(7, 7),
|
|
strides=2,
|
|
padding="same")
|
|
self.bn1 = tf.keras.layers.BatchNormalization()
|
|
self.pool1 = tf.keras.layers.MaxPool2D(pool_size=(3, 3),
|
|
strides=2,
|
|
padding="same")
|
|
else:
|
|
self.conv1 = tf.keras.layers.Conv2DTranspose(filters=3 if labels is None else labels,
|
|
kernel_size=(7, 7),
|
|
strides=2 if len(layer_params) > 1 else 1,
|
|
padding="same")
|
|
self.bn1 = tf.keras.layers.BatchNormalization()
|
|
self.pool1 = tf.keras.layers.UpSampling2D()
|
|
|
|
self.layer1 = make_basic_block_layer(filter_num=32 if len(layer_params) > 1 else (features if not is_cvae else features*2),
|
|
blocks=layer_params[0],
|
|
stride=1 if not reversed else 2,
|
|
upsample=reversed)
|
|
if len(layer_params) >= 2:
|
|
self.layer2 = make_basic_block_layer(filter_num=32,
|
|
blocks=layer_params[1],
|
|
stride=2,
|
|
upsample=reversed)
|
|
if len(layer_params) >= 3:
|
|
self.layer3 = make_basic_block_layer(filter_num=32,
|
|
blocks=layer_params[2],
|
|
stride=2,
|
|
upsample=reversed)
|
|
if len(layer_params) >= 4:
|
|
self.layer4 = make_basic_block_layer(filter_num=64,
|
|
blocks=layer_params[3],
|
|
stride=2,
|
|
upsample=reversed)
|
|
|
|
if len(layer_params) >= 5:
|
|
self.layer5 = make_basic_block_layer(filter_num=features if not is_cvae else features*2,
|
|
blocks=layer_params[4],
|
|
stride=2 if not reversed else 1,
|
|
upsample=reversed)
|
|
|
|
def call(self, inputs, training=None, mask=None):
|
|
if not self.reversed:
|
|
x = self.conv1(inputs)
|
|
x = self.bn1(x, training=training)
|
|
x = tf.nn.relu(x)
|
|
x = self.pool1(x)
|
|
x = self.layer1(x, training=training)
|
|
if len(self.layer_params) >= 2:
|
|
x = self.layer2(x, training=training)
|
|
if len(self.layer_params) >= 3:
|
|
x = self.layer3(x, training=training)
|
|
if len(self.layer_params) >= 4:
|
|
x = self.layer4(x, training=training)
|
|
if len(self.layer_params) >= 5:
|
|
x = self.layer5(x, training=training)
|
|
else:
|
|
x = inputs
|
|
if len(self.layer_params) >= 5:
|
|
x = self.layer5(x, training=training)
|
|
if len(self.layer_params) >= 4:
|
|
x = self.layer4(x, training=training)
|
|
if len(self.layer_params) >= 2:
|
|
x = self.layer3(x, training=training)
|
|
if len(self.layer_params) >= 3:
|
|
x = self.layer2(x, training=training)
|
|
x = self.layer1(x, training=training)
|
|
|
|
x = self.pool1(x)
|
|
x = self.bn1(x, training=training)
|
|
# x = tf.nn.relu(x)
|
|
x = self.conv1(x)
|
|
if self.labels is None:
|
|
x = tf.nn.tanh(x)
|
|
# else:
|
|
# x = tf.nn.softmax(x)
|
|
|
|
return x
|
|
|
|
|
|
class ResNetTypeII(tf.keras.Model):
|
|
def __init__(self, layer_params, features, reversed=False, labels=None, is_cvae=False):
|
|
super(ResNetTypeII, self).__init__()
|
|
self.reversed = reversed
|
|
self.labels = labels
|
|
|
|
if not reversed:
|
|
self.conv1 = tf.keras.layers.Conv2D(filters=32,
|
|
kernel_size=(7, 7),
|
|
strides=2,
|
|
padding="same")
|
|
self.bn1 = tf.keras.layers.BatchNormalization()
|
|
self.pool1 = tf.keras.layers.MaxPool2D(pool_size=(3, 3),
|
|
strides=2,
|
|
padding="same")
|
|
else:
|
|
self.conv1 = tf.keras.layers.Conv2DTranspose(filters=3 if labels is None else labels,
|
|
kernel_size=(7, 7),
|
|
strides=2,
|
|
padding="same")
|
|
self.bn1 = tf.keras.layers.BatchNormalization()
|
|
self.pool1 = tf.keras.layers.UpSampling2D()
|
|
|
|
self.layer1 = make_bottleneck_layer(filter_num=32,
|
|
blocks=layer_params[0],
|
|
stride=1 if not reversed else 2,
|
|
upsample=reversed)
|
|
self.layer2 = make_bottleneck_layer(filter_num=64,
|
|
blocks=layer_params[1],
|
|
stride=2,
|
|
upsample=reversed)
|
|
self.layer3 = make_bottleneck_layer(filter_num=64,
|
|
blocks=layer_params[2],
|
|
stride=2,
|
|
upsample=reversed)
|
|
self.layer4 = make_bottleneck_layer(filter_num=128,
|
|
blocks=layer_params[3],
|
|
stride=2,
|
|
upsample=reversed)
|
|
self.layer5 = make_bottleneck_layer(filter_num=features if not is_cvae else features*2,
|
|
blocks=layer_params[4],
|
|
stride=2 if not reversed else 1,
|
|
upsample=reversed)
|
|
|
|
|
|
def call(self, inputs, training=None, mask=None):
|
|
if not self.reversed:
|
|
x = self.conv1(inputs)
|
|
x = self.bn1(x, training=training)
|
|
x = tf.nn.relu(x)
|
|
x = self.pool1(x)
|
|
x = self.layer1(x, training=training)
|
|
x = self.layer2(x, training=training)
|
|
x = self.layer3(x, training=training)
|
|
x = self.layer4(x, training=training)
|
|
x = self.layer5(x, training=training)
|
|
else:
|
|
x = self.layer5(inputs, training=training)
|
|
x = self.layer4(x, training=training)
|
|
x = self.layer3(x, training=training)
|
|
x = self.layer2(x, training=training)
|
|
x = self.layer1(x, training=training)
|
|
|
|
x = self.pool1(x)
|
|
x = self.bn1(x, training=training)
|
|
x = self.conv1(x)
|
|
if self.labels is None:
|
|
x = tf.nn.tanh(x)
|
|
|
|
return x
|
|
|
|
|
|
def resnet_18(layer_params, features, reversed=False, labels=None, is_cvae=False):
|
|
return ResNetTypeI(layer_params=layer_params, features=features, reversed=reversed, labels=labels, is_cvae=is_cvae)
|
|
|
|
|
|
def resnet_34():
|
|
return ResNetTypeI(layer_params=[3, 4, 6, 3])
|
|
|
|
|
|
def resnet_50():
|
|
return ResNetTypeII(layer_params=[3, 4, 6, 3])
|
|
|
|
|
|
def resnet_101(layer_params, features, reversed=False, labels=None, is_cvae=False):
|
|
return ResNetTypeII(layer_params=layer_params, features=features, reversed=reversed, labels=labels, is_cvae=is_cvae)
|
|
|
|
|
|
def resnet_152():
|
|
return ResNetTypeII(layer_params=[3, 8, 36, 3]) |