2martens/masterthesis
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Added ability to load trained weights

Browse Source 
Signed-off-by: Jim Martens <github@2martens.de>
This commit is contained in:
Jim Martens
2019-03-20 14:54:28 +01:00
parent 6ac3b5fdcf
commit 9a7a906acf
1 changed files with 14 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/twomartens/masterthesis/ssd.py
View File

@ -29,6 +29,8 @@ Classes:
``SSD``: wraps vanilla SSD 300 model
"""
from typing import Optional
import tensorflow as tf
from twomartens.masterthesis.ssd_keras.models import keras_ssd300
@ -47,12 +49,17 @@ class SSD:
Args:
mode: one of training, inference, and inference_fast
weights_path: path to trained weights
"""
def __init__(self, mode: str) -> None:
def __init__(self, mode: str, weights_path: Optional[str] = None) -> None:
self._model = keras_ssd300.ssd_300(image_size=IMAGE_SIZE, n_classes=N_CLASSES, mode=mode)
self.mode = mode
# load existing weights
if weights_path is not None:
self._model.load_weights(weights_path, by_name=True)
def __call__(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor:
return self._model(inputs)
@ -63,12 +70,17 @@ class DropoutSSD:
Args:
mode: one of training, inference, and inference_fast
weights_path: path to trained weights
"""
def __init__(self, mode: str) -> None:
def __init__(self, mode: str, weights_path: Optional[str] = None) -> None:
self._model = keras_ssd300_dropout.ssd_300_dropout(image_size=IMAGE_SIZE, n_classes=N_CLASSES,
dropout_rate=DROPOUT_RATE, mode=mode)
self.mode = mode
# load existing weights
if weights_path is not None:
self._model.load_weights(weights_path, by_name=True)
def __call__(self, inputs: tf.Tensor, *args, **kwargs) -> tf.Tensor:
return self._model(inputs)