Add documentation and README, use Upsampling2D rather than image Resizing layer
This commit is contained in:
Piv
@@ -2,10 +2,22 @@ import tensorflow as tf
|
||||
import tensorflow.keras as keras
|
||||
import tensorflow_datasets as tfds
|
||||
|
||||
"""
|
||||
Unofficial tensorflow keras implementation of FastDepth (mobilenet_nnconv5).
|
||||
PyTorch (official) Fast Depth Implementation: https://github.com/dwofk/fast-depth
|
||||
|
||||
# Ripped from: https://forums.developer.nvidia.com/t/could-not-create-cudnn-handle-cudnn-status-alloc-failed/108261/4?u=mpivato4
|
||||
There's also an experimental version that does not use BatchNormalisation, as well as Parametric ReLU and bilinear
|
||||
upsampling (mobilenet_nnconv5_no_bn)
|
||||
"""
|
||||
|
||||
|
||||
# Ripped from:
|
||||
# https://forums.developer.nvidia.com/t/could-not-create-cudnn-handle-cudnn-status-alloc-failed/108261/4?u=mpivato4
|
||||
# Seems to be an issue on windows so explicitly set gpu growth
|
||||
def fix_windows_gpu():
|
||||
"""
|
||||
Fixes Windows GPU bug when attempting to allocate memory using cuDNN
|
||||
"""
|
||||
gpus = tf.config.experimental.list_physical_devices('GPU')
|
||||
if gpus:
|
||||
try:
|
||||
@@ -19,11 +31,6 @@ def fix_windows_gpu():
|
||||
print(e)
|
||||
|
||||
|
||||
'''
|
||||
Functional version of fastdepth model
|
||||
'''
|
||||
|
||||
|
||||
def FDDepthwiseBlock(inputs,
|
||||
out_channels,
|
||||
block_id=1):
|
||||
@@ -35,14 +42,13 @@ def FDDepthwiseBlock(inputs,
|
||||
return keras.layers.ReLU(6., name='conv_pw_%d_relu' % block_id)(x)
|
||||
|
||||
|
||||
def FDDepthwiseBlockNoBN(inputs, out_channels, block_id=1):
|
||||
x = keras.layers.DepthwiseConv2D(5, padding='same')(inputs)
|
||||
x = keras.layers.ReLU(6.)(x)
|
||||
x = keras.layers.Conv2D(out_channels, 1, padding='same')(x)
|
||||
return keras.layers.ReLU(6., name='conv_pw_%d_relu' % block_id)(x)
|
||||
|
||||
|
||||
def make_mobilenet_nnconv5(weights=None, shape=(224, 224, 3)):
|
||||
def mobilenet_nnconv5(weights=None, shape=(224, 224, 3)):
|
||||
"""
|
||||
Replication of the FastDepth model in Tensorflow, using the keras Functional API
|
||||
:param weights: Pretrained weights for MobileNet, defaults to None
|
||||
:param shape: Input shape of the image, defaults to (224, 224, 3)
|
||||
:return: FastDepth keras Model
|
||||
"""
|
||||
input = keras.layers.Input(shape=shape)
|
||||
mobilenet = keras.applications.MobileNet(input_tensor=input, include_top=False, weights=weights)
|
||||
for layer in mobilenet.layers:
|
||||
@@ -50,21 +56,19 @@ def make_mobilenet_nnconv5(weights=None, shape=(224, 224, 3)):
|
||||
|
||||
# Fast depth decoder
|
||||
x = FDDepthwiseBlock(mobilenet.output, 512, block_id=14)
|
||||
# TODO: Bilinear interpolation
|
||||
# x = keras.layers.experimental.preprocessing.Resizing(14, 14, interpolation='bilinear')
|
||||
# Nearest neighbour interpolation, used by fast depth paper
|
||||
x = keras.layers.experimental.preprocessing.Resizing(14, 14, interpolation='nearest')(x)
|
||||
x = keras.layers.UpSampling2D()(x)
|
||||
x = FDDepthwiseBlock(x, 256, block_id=15)
|
||||
x = keras.layers.experimental.preprocessing.Resizing(28, 28, interpolation='nearest')(x)
|
||||
x = keras.layers.UpSampling2D()(x)
|
||||
x = keras.layers.Add()([x, mobilenet.get_layer(name="conv_pw_5_relu").output])
|
||||
x = FDDepthwiseBlock(x, 128, block_id=16)
|
||||
x = keras.layers.experimental.preprocessing.Resizing(56, 56, interpolation='nearest')(x)
|
||||
x = keras.layers.UpSampling2D()(x)
|
||||
x = keras.layers.Add()([x, mobilenet.get_layer(name="conv_pw_3_relu").output])
|
||||
x = FDDepthwiseBlock(x, 64, block_id=17)
|
||||
x = keras.layers.experimental.preprocessing.Resizing(112, 112, interpolation='nearest')(x)
|
||||
x = keras.layers.UpSampling2D()(x)
|
||||
x = keras.layers.Add()([x, mobilenet.get_layer(name="conv_pw_1_relu").output])
|
||||
x = FDDepthwiseBlock(x, 32, block_id=18)
|
||||
x = keras.layers.experimental.preprocessing.Resizing(224, 224, interpolation='nearest')(x)
|
||||
x = keras.layers.UpSampling2D()(x)
|
||||
|
||||
x = keras.layers.Conv2D(1, 1, padding='same')(x)
|
||||
x = keras.layers.BatchNormalization()(x)
|
||||
@@ -72,7 +76,25 @@ def make_mobilenet_nnconv5(weights=None, shape=(224, 224, 3)):
|
||||
return keras.Model(inputs=input, outputs=x, name="fast_depth")
|
||||
|
||||
|
||||
def make_mobilenet_nnconv5_no_bn(weights=None, shape=(224, 224, 3)):
|
||||
#### Experimental ####
|
||||
def FDDepthwiseBlockNoBN(inputs, out_channels, block_id=1):
|
||||
x = keras.layers.DepthwiseConv2D(5, padding='same')(inputs)
|
||||
x = keras.layers.PReLU()(x)
|
||||
x = keras.layers.Conv2D(out_channels, 1, padding='same')(x)
|
||||
return keras.layers.PReLU(name='conv_pw_%d_relu' % block_id)(x)
|
||||
|
||||
|
||||
def mobilenet_nnconv5_no_bn(weights=None, shape=(224, 224, 3)):
|
||||
"""
|
||||
Experimental version of the FastDepth model.
|
||||
This version has the following changes:
|
||||
- Bilinear upsampling is used rather than nearest neighbour
|
||||
- No BatchNormalisation in decoder
|
||||
- Parametric ReLU in Decoder rather than ReLU
|
||||
:param weights: Pretrained weights for MobileNet, defaults to None
|
||||
:param shape: Input shape of the image, defaults to (224, 224, 3)
|
||||
:return: Experimental FastDepth keras Model
|
||||
"""
|
||||
input = keras.layers.Input(shape=shape)
|
||||
mobilenet = keras.applications.MobileNet(input_tensor=input, include_top=False, weights=weights)
|
||||
for layer in mobilenet.layers:
|
||||
@@ -80,26 +102,24 @@ def make_mobilenet_nnconv5_no_bn(weights=None, shape=(224, 224, 3)):
|
||||
|
||||
# Fast depth decoder
|
||||
x = FDDepthwiseBlockNoBN(mobilenet.output, 512, block_id=14)
|
||||
# Nearest neighbour interpolation, used by fast depth paper
|
||||
x = keras.layers.experimental.preprocessing.Resizing(14, 14, interpolation='bilinear')(x)
|
||||
x = keras.layers.UpSampling2D(interpolation='bilinear')(x)
|
||||
x = FDDepthwiseBlockNoBN(x, 256, block_id=15)
|
||||
x = keras.layers.experimental.preprocessing.Resizing(28, 28, interpolation='bilinear')(x)
|
||||
x = keras.layers.UpSampling2D(interpolation='bilinear')(x)
|
||||
x = keras.layers.Add()([x, mobilenet.get_layer(name="conv_pw_5_relu").output])
|
||||
x = FDDepthwiseBlockNoBN(x, 128, block_id=16)
|
||||
x = keras.layers.experimental.preprocessing.Resizing(56, 56, interpolation='bilinear')(x)
|
||||
x = keras.layers.UpSampling2D(interpolation='bilinear')(x)
|
||||
x = keras.layers.Add()([x, mobilenet.get_layer(name="conv_pw_3_relu").output])
|
||||
x = FDDepthwiseBlockNoBN(x, 64, block_id=17)
|
||||
x = keras.layers.experimental.preprocessing.Resizing(112, 112, interpolation='bilinear')(x)
|
||||
x = keras.layers.UpSampling2D(interpolation='bilinear')(x)
|
||||
x = keras.layers.Add()([x, mobilenet.get_layer(name="conv_pw_1_relu").output])
|
||||
x = FDDepthwiseBlockNoBN(x, 32, block_id=18)
|
||||
x = keras.layers.experimental.preprocessing.Resizing(224, 224, interpolation='bilinear')(x)
|
||||
x = keras.layers.UpSampling2D(interpolation='bilinear')(x)
|
||||
|
||||
x = keras.layers.Conv2D(1, 1, padding='same')(x)
|
||||
x = keras.layers.ReLU(6.)(x)
|
||||
return keras.Model(inputs=input, outputs=x, name="fast_depth")
|
||||
x = keras.layers.PReLU()(x)
|
||||
return keras.Model(inputs=input, outputs=x, name="fast_depth_experimental")
|
||||
|
||||
|
||||
# TODO: Fix these, float doesn't work same as pytorch
|
||||
def delta1_metric(y_true, y_pred):
|
||||
maxRatio = tf.maximum(y_pred / y_true, y_true / y_pred)
|
||||
return tf.nn.moments(tf.cast(maxRatio < tf.convert_to_tensor(1.25), tf.float32), axes=None)[0]
|
||||
@@ -116,6 +136,10 @@ def delta3_metric(y_true, y_pred):
|
||||
|
||||
|
||||
def compile(model):
|
||||
"""
|
||||
Compile FastDepth model with relevant metrics
|
||||
:param model: Model to compile
|
||||
"""
|
||||
# TODO: Learning rate (exponential decay)
|
||||
model.compile(optimizer=keras.optimizers.SGD(momentum=0.9),
|
||||
loss=keras.losses.MeanSquaredError(),
|
||||
@@ -127,8 +151,17 @@ def compile(model):
|
||||
|
||||
|
||||
def train(existing_model=None, pretrained_weights='imagenet', epochs=4, save_file=None, dataset=None):
|
||||
"""
|
||||
Compile, train and save (if a save file is specified) a Fast Depth model.
|
||||
:param existing_model: Existing FastDepth model to train. None will create
|
||||
:param pretrained_weights: Weights to use if existing_model is not specified. See keras.applications.MobileNet
|
||||
weights parameter for options here.
|
||||
:param epochs: Number of epochs to run for
|
||||
:param save_file: File/directory to save to after training. By default the model won't be saved
|
||||
:param dataset: Train dataset to use. By default will DOWNLOAD and use tensorflow nyu_v2 dataset
|
||||
"""
|
||||
if not existing_model:
|
||||
existing_model = make_mobilenet_nnconv5(pretrained_weights)
|
||||
existing_model = mobilenet_nnconv5(pretrained_weights)
|
||||
compile(existing_model)
|
||||
if not dataset:
|
||||
dataset = load_nyu()
|
||||
@@ -162,6 +195,11 @@ def forward(model, image):
|
||||
|
||||
|
||||
def load_model(file):
|
||||
"""
|
||||
Load previously trained FastDepth model from disk. Will include relevant metrics (custom objects)
|
||||
:param file: File/directory to load the model from
|
||||
:return:
|
||||
"""
|
||||
return keras.models.load_model(file, custom_objects={'delta1_metric': delta1_metric,
|
||||
'delta2_metric': delta2_metric,
|
||||
'delta3_metric': delta3_metric})
|
||||
@@ -181,6 +219,10 @@ def crop_and_resize(x):
|
||||
|
||||
|
||||
def load_nyu():
|
||||
"""
|
||||
Load the nyu_v2 dataset train split. Will be downloaded to ../nyu
|
||||
:return: nyu_v2 dataset builder
|
||||
"""
|
||||
builder = tfds.builder('nyu_depth_v2')
|
||||
builder.download_and_prepare(download_dir='../nyu')
|
||||
return builder \
|
||||
@@ -191,6 +233,10 @@ def load_nyu():
|
||||
|
||||
|
||||
def load_nyu_evaluate():
|
||||
"""
|
||||
Load the nyu_v2 dataset validation split. Will be downloaded to ../nyu
|
||||
:return: nyu_v2 dataset builder
|
||||
"""
|
||||
builder = tfds.builder('nyu_depth_v2')
|
||||
builder.download_and_prepare(download_dir='../nyu')
|
||||
return builder.as_dataset(split='validation').batch(1).map(lambda x: crop_and_resize(x))
|
||||
|
||||
Reference in New Issue
Block a user