Commit cb469550 authored by Laurent COLBOIS's avatar Laurent COLBOIS

Merge branch 'mxnet-demanding' into 'master'

Resolve "ArcFace (MXNET) + IJBC lead to memory error"

Closes #42

See merge request !104
parents 1993f551 71387a13
Pipeline #48375 passed with stages
in 55 minutes and 52 seconds
......@@ -9,6 +9,9 @@ from bob.bio.base.pipelines.vanilla_biometrics import (
if "database" in locals():
annotation_type = database.annotation_type
fixed_positions = database.fixed_positions
memory_demanding = (
database.memory_demanding if hasattr(database, "memory_demanding") else False
)
else:
annotation_type = None
fixed_positions = None
......@@ -16,7 +19,7 @@ else:
def load(annotation_type, fixed_positions=None):
transformer = embedding_transformer_112x112(
ArcFaceInsightFace(), annotation_type, fixed_positions, color_channel="rgb"
ArcFaceInsightFace(memory_demanding=memory_demanding), annotation_type, fixed_positions, color_channel="rgb"
)
algorithm = Distance()
......
......@@ -21,10 +21,11 @@ class ArcFaceInsightFace(TransformerMixin, BaseEstimator):
"""
def __init__(self, use_gpu=False, **kwargs):
def __init__(self, use_gpu=False, memory_demanding=False, **kwargs):
super().__init__(**kwargs)
self.model = None
self.use_gpu = use_gpu
self.memory_demanding = memory_demanding
internal_path = pkg_resources.resource_filename(
__name__, os.path.join("data", "arcface_insightface"),
......@@ -76,10 +77,17 @@ class ArcFaceInsightFace(TransformerMixin, BaseEstimator):
self.load_model()
X = check_array(X, allow_nd=True)
X = mx.nd.array(X)
db = mx.io.DataBatch(data=(X,))
self.model.forward(db, is_train=False)
return self.model.get_outputs()[0].asnumpy()
def _transform(X):
X = mx.nd.array(X)
db = mx.io.DataBatch(data=(X,))
self.model.forward(db, is_train=False)
return self.model.get_outputs()[0].asnumpy()
if self.memory_demanding:
return np.array([_transform(x[None, ...]) for x in X])
else:
return _transform(X)
def __getstate__(self):
# Handling unpicklable objects
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment