|
import pytest |
|
|
|
from easydict import EasyDict |
|
from ding.framework import OnlineRLContext |
|
from ding.framework.middleware.ckpt_handler import CkptSaver |
|
|
|
import torch.nn as nn |
|
import torch.optim as optim |
|
import os |
|
import shutil |
|
|
|
from unittest.mock import Mock, patch |
|
from ding.framework import task |
|
from ding.policy.base_policy import Policy |
|
|
|
|
|
class TheModelClass(nn.Module): |
|
|
|
def state_dict(self): |
|
return 'fake_state_dict' |
|
|
|
|
|
class MockPolicy(Mock): |
|
|
|
def __init__(self, model, **kwargs) -> None: |
|
super(MockPolicy, self).__init__(model) |
|
self.learn_mode = model |
|
|
|
@property |
|
def eval_mode(self): |
|
return EasyDict({"state_dict": lambda: {}}) |
|
|
|
|
|
@pytest.mark.unittest |
|
def test_ckpt_saver(): |
|
exp_name = 'test_ckpt_saver_exp' |
|
|
|
ctx = OnlineRLContext() |
|
|
|
train_freq = 100 |
|
model = TheModelClass() |
|
|
|
if not os.path.exists(exp_name): |
|
os.makedirs(exp_name) |
|
|
|
prefix = '{}/ckpt'.format(exp_name) |
|
|
|
with patch("ding.policy.Policy", MockPolicy), task.start(): |
|
policy = MockPolicy(model) |
|
|
|
def mock_save_file(path, data, fs_type=None, use_lock=False): |
|
assert path == "{}/eval.pth.tar".format(prefix) |
|
|
|
with patch("ding.framework.middleware.ckpt_handler.save_file", mock_save_file): |
|
ctx.train_iter = 1 |
|
ctx.eval_value = 9.4 |
|
ckpt_saver = CkptSaver(policy, exp_name, train_freq) |
|
ckpt_saver(ctx) |
|
|
|
def mock_save_file(path, data, fs_type=None, use_lock=False): |
|
assert path == "{}/iteration_{}.pth.tar".format(prefix, ctx.train_iter) |
|
|
|
with patch("ding.framework.middleware.ckpt_handler.save_file", mock_save_file): |
|
ctx.train_iter = 100 |
|
ctx.eval_value = 1 |
|
ckpt_saver(ctx) |
|
|
|
def mock_save_file(path, data, fs_type=None, use_lock=False): |
|
assert path == "{}/final.pth.tar".format(prefix) |
|
|
|
with patch("ding.framework.middleware.ckpt_handler.save_file", mock_save_file): |
|
task.finish = True |
|
ckpt_saver(ctx) |
|
|
|
shutil.rmtree(exp_name) |
|
|
