MTCNN is not serializable
View options
Closed Issue created
With the current implementation, MTCNN is not serializable via pickle. When used locally, it is not an issue, however, when running on the cluster, we are getting lovely messages into workers' logs such as:
distributed.protocol.pickle - INFO - Failed to serialize CheckpointWrapper(estimator=SampleWrapper(estimator=BoundingBoxAnnotatorCrop(annotator=MTCNN(thresholds=(0.1,
0.2,
0.2)),
eyes_cropper=FaceEyesNorm(final_image_size=(112,
112),
reference_eyes_location={'bottomright': (112,
112),
'leye': (55,
72),
'reye': (55,
40),
'topleft': (0,
0)})),
fit_extra_arguments=(),
input_attribute='data',
output_attribute='data',
transform_...',
'annotations'),)),
extension='.h5',
features_dir='/idiap/temp/cecabert/experiments/bob101/results-sge/ijbc/cropper',
hash_fn=<function hash_string at 0x7fa50d891000>,
load_func=<function load at 0x7fa507f5ed40>,
model_path='/idiap/temp/cecabert/experiments/bob101/results-sge/ijbc/cropper.pkl',
sample_attribute='data',
save_func=<function save at 0x7fa507f5ef80>). Exception: cannot pickle '_thread.RLock' object
distributed.protocol.core - CRITICAL - Failed to deserialize
Traceback (most recent call last):
File "/remote/idiap.svm/temp.biometric03/cecabert/mambaforge/envs/bob_deps/lib/python3.10/site-packages/distributed/protocol/core.py", line 111, in loads
return msgpack.loads(
File "msgpack/_unpacker.pyx", line 194, in msgpack._cmsgpack.unpackb
File "/remote/idiap.svm/temp.biometric03/cecabert/mambaforge/envs/bob_deps/lib/python3.10/site-packages/distributed/protocol/core.py", line 103, in _decode_default
return merge_and_deserialize(
File "/remote/idiap.svm/temp.biometric03/cecabert/mambaforge/envs/bob_deps/lib/python3.10/site-packages/distributed/protocol/serialize.py", line 488, in merge_and_deserialize
return deserialize(header, merged_frames, deserializers=deserializers)
File "/remote/idiap.svm/temp.biometric03/cecabert/mambaforge/envs/bob_deps/lib/python3.10/site-packages/distributed/protocol/serialize.py", line 417, in deserialize
return loads(header, frames)
File "/remote/idiap.svm/temp.biometric03/cecabert/mambaforge/envs/bob_deps/lib/python3.10/site-packages/distributed/protocol/serialize.py", line 180, in serialization_error_loads
raise TypeError(msg)
TypeError: Could not serialize object of type CheckpointWrapper.
Traceback (most recent call last):
File "/remote/idiap.svm/temp.biometric03/cecabert/mambaforge/envs/bob_deps/lib/python3.10/site-packages/distributed/protocol/pickle.py", line 40, in dumps
result = pickle.dumps(x, **dump_kwargs)
AttributeError: Can't pickle local object 'WeakSet.__init__.<locals>._remove'The issue it that in some case, it tries to serialize the already loaded underlying Tensorflow Graph. This can be solved with the same mechanism used in PyTorchModel class by overriding the __getstate__ method as follow:
def __getstate__(self):
# Handling unpicklable objects
state = {}
for key, value in super().__getstate__().items():
if key != '_fun':
state[key] = value
state['_fun'] = None
return stateWith this change, the serialization now works properly and can be tested with:
mtcnn = MTCNN()
mtcnn.mtcnn_fun # Force instantiation of TF graph
other = pickle.loads(pickle.dumps(mtcnn))
# No AttributeError: Can't pickle local object 'WeakSet.__init__.<locals>._remove'
# TF graph will be lazily initialized if needed