mirror of
https://github.com/gryf/coach.git
synced 2025-12-17 19:20:19 +01:00
67eb9e4c28
* Adding checkpointing framework as well as mxnet checkpointing implementation. - MXNet checkpoint for each network is saved in a separate file. * Adding checkpoint restore for mxnet to graph-manager * Add unit-test for get_checkpoint_state() * Added match.group() to fix unit-test failing on CI * Added ONNX export support for MXNet
22 lines
973 B
Python
22 lines
973 B
Python
import pytest
|
|
|
|
from rl_coach import utils
|
|
|
|
|
|
@pytest.mark.unit_test
|
|
def test_get_checkpoint_state_default():
|
|
files = ['4.test.ckpt.ext', '2.test.ckpt.ext', '3.test.ckpt.ext', '1.test.ckpt.ext']
|
|
checkpoint_state = utils.get_checkpoint_state(files)
|
|
assert checkpoint_state.model_checkpoint_path == '4.test.ckpt'
|
|
assert checkpoint_state.all_model_checkpoint_paths == [f[:-4] for f in sorted(files)]
|
|
|
|
|
|
@pytest.mark.unit_test
|
|
def test_get_checkpoint_state_custom():
|
|
files = ['prefix.4.test.ckpt.ext', 'prefix.2.test.ckpt.ext', 'prefix.3.test.ckpt.ext', 'prefix.1.test.ckpt.ext']
|
|
assert len(utils.get_checkpoint_state(files).all_model_checkpoint_paths) == 0 # doesn't match the default pattern
|
|
checkpoint_state = utils.get_checkpoint_state(files, filename_pattern=r'([0-9]+)[^0-9].*?\.ckpt')
|
|
assert checkpoint_state.model_checkpoint_path == '4.test.ckpt'
|
|
assert checkpoint_state.all_model_checkpoint_paths == [f[7:-4] for f in sorted(files)]
|
|
|