Files
blackbird/actpred/test_actpred.py
T
qing 3d03f7a1ce actpred: 代码审计(16项发现全部修复)+ 文档更新
审计报告: docs/reviews/actpred_CODE_AUDIT_20260808.md
- P0(2项): episode索引错位(dataset.py)、推理归一化静默失败(predict.py)
- P1(6项): 空训练集校验、图像自然排序、空样本防御、形状校验、
  MobileNet下载错误提示、stats路径校验
- P2(8项): README示例修正、公共属性、glob优化、版本约束、
  单元测试(17项)、GPU安装说明、BooleanOptionalAction、ONNX验证

新增文件: actpred/test_actpred.py (17项单元测试, 全部通过)
更新文档: README.md配置表扩充、INDEX.md、审计报告修复总结
2026-08-08 16:05:01 +00:00

220 lines
8.3 KiB
Python

"""actpred 模块单元测试
运行方式:
cd actpred && python -m pytest test_actpred.py -v
# 或
cd actpred && python test_actpred.py
"""
import os
import tempfile
import unittest
from pathlib import Path
import numpy as np
from PIL import Image
import torch
from config import Config, DataConfig, ModelConfig, build_argparser, apply_overrides
from dataset import _natural_key, ActionDataset, DataRecorder
from model import ActionPredictor
from predict import parse_action_history
class TestNaturalKey(unittest.TestCase):
"""测试自然排序键函数"""
def test_natural_order_no_padding(self):
"""无前导零的文件名应按数字序排列"""
files = [Path(f"{i}.jpg") for i in [3, 1, 10, 2, 20]]
sorted_files = sorted(files, key=_natural_key)
names = [f.name for f in sorted_files]
self.assertEqual(names, ["1.jpg", "2.jpg", "3.jpg", "10.jpg", "20.jpg"])
def test_natural_order_with_padding(self):
"""有前导零的文件名排序也应正确"""
files = [Path(f"{i:04d}.jpg") for i in [3, 1, 10, 2, 20]]
sorted_files = sorted(files, key=_natural_key)
names = [f.name for f in sorted_files]
self.assertEqual(names, ["0001.jpg", "0002.jpg", "0003.jpg", "0010.jpg", "0020.jpg"])
class TestParseActionHistory(unittest.TestCase):
"""测试动作历史解析函数"""
def test_single_step(self):
history = parse_action_history("1,2,3,4", action_dim=4)
self.assertEqual(history.shape, (1, 4))
np.testing.assert_array_equal(history[0], [1, 2, 3, 4])
def test_multi_step(self):
history = parse_action_history("1,0,0,0:0,1,0,0:0,0,1,0", action_dim=4)
self.assertEqual(history.shape, (3, 4))
def test_wrong_dim_raises(self):
with self.assertRaises(AssertionError):
parse_action_history("1,2,3", action_dim=4)
class TestActionPredictor(unittest.TestCase):
"""测试模型前向传播"""
def test_forward_simple_mlp(self):
data_cfg = DataConfig(image_size=64, action_dim=8, history_len=5)
model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp")
model = ActionPredictor(model_cfg, data_cfg)
model.eval()
image = torch.randn(2, 3, 64, 64)
history = torch.randn(2, 5, 8)
with torch.no_grad():
out = model(image, history)
self.assertEqual(out.shape, (2, 8))
def test_forward_simple_gru(self):
data_cfg = DataConfig(image_size=64, action_dim=8, history_len=5)
model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="gru")
model = ActionPredictor(model_cfg, data_cfg)
model.eval()
image = torch.randn(2, 3, 64, 64)
history = torch.randn(2, 5, 8)
with torch.no_grad():
out = model(image, history)
self.assertEqual(out.shape, (2, 8))
def test_forward_dynamic_batch(self):
data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3)
model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp")
model = ActionPredictor(model_cfg, data_cfg)
model.eval()
for batch_size in [1, 4, 8]:
image = torch.randn(batch_size, 3, 32, 32)
history = torch.randn(batch_size, 3, 4)
with torch.no_grad():
out = model(image, history)
self.assertEqual(out.shape, (batch_size, 4))
class TestDataRecorder(unittest.TestCase):
"""测试数据录制器"""
def test_record_and_flush(self):
with tempfile.TemporaryDirectory() as tmpdir:
recorder = DataRecorder(tmpdir, action_dim=4)
for _ in range(10):
img = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)
action = np.random.randn(4).astype(np.float32)
recorder.record(img, action)
episode_path = recorder.flush()
self.assertTrue(os.path.exists(os.path.join(episode_path, "actions.npy")))
images_dir = os.path.join(episode_path, "images")
self.assertTrue(os.path.isdir(images_dir))
self.assertEqual(len(os.listdir(images_dir)), 10)
actions = np.load(os.path.join(episode_path, "actions.npy"))
self.assertEqual(actions.shape, (10, 4))
def test_multiple_episodes_increment_index(self):
with tempfile.TemporaryDirectory() as tmpdir:
r1 = DataRecorder(tmpdir, action_dim=4)
r1.flush()
r2 = DataRecorder(tmpdir, action_dim=4)
r2.flush()
self.assertNotEqual(r1.episode_dir, r2.episode_dir)
self.assertIn("episode_0001", str(r2.episode_dir))
class TestActionDataset(unittest.TestCase):
"""测试数据集加载"""
@staticmethod
def _create_dummy_episode(base_dir, n_frames=10, action_dim=4, image_size=32):
ep_dir = base_dir / "episode_0001"
images_dir = ep_dir / "images"
images_dir.mkdir(parents=True, exist_ok=True)
actions = []
for i in range(n_frames):
img = Image.fromarray(
np.random.randint(0, 255, (image_size, image_size, 3), dtype=np.uint8))
img.save(images_dir / f"{i:04d}.jpg")
actions.append(np.random.randn(action_dim).astype(np.float32))
np.save(ep_dir / "actions.npy", np.array(actions))
def test_dataset_loading(self):
with tempfile.TemporaryDirectory() as tmpdir:
self._create_dummy_episode(Path(tmpdir), n_frames=10, action_dim=4)
ds = ActionDataset(tmpdir, history_len=3, action_dim=4, image_size=32)
self.assertEqual(len(ds), 7) # 10 - 3 = 7 samples
sample = ds[0]
self.assertIn("image", sample)
self.assertIn("action_history", sample)
self.assertIn("target", sample)
self.assertEqual(sample["image"].shape, (3, 32, 32))
self.assertEqual(sample["action_history"].shape, (3, 4))
self.assertEqual(sample["target"].shape, (4,))
def test_episode_count(self):
with tempfile.TemporaryDirectory() as tmpdir:
self._create_dummy_episode(Path(tmpdir))
count = ActionDataset.get_episode_count(tmpdir)
self.assertEqual(count, 1)
def test_empty_stats_raises(self):
"""空数据集计算归一化统计量应报错"""
with tempfile.TemporaryDirectory() as tmpdir:
with self.assertRaises(ValueError):
ActionDataset(tmpdir, history_len=5, action_dim=4,
normalize_actions=True)
def test_invalid_stats_path_raises(self):
"""指定的 stats 路径不存在应报错"""
with tempfile.TemporaryDirectory() as tmpdir:
with self.assertRaises(FileNotFoundError):
ActionDataset(tmpdir, history_len=3, action_dim=4,
normalize_actions=True,
action_stats_path="/nonexistent/path.npy")
class TestConfig(unittest.TestCase):
"""测试配置和 CLI 参数"""
def test_default_config(self):
cfg = Config()
self.assertEqual(cfg.data.image_size, 64)
self.assertEqual(cfg.data.action_dim, 8)
self.assertEqual(cfg.data.history_len, 5)
self.assertEqual(cfg.model.cnn_backbone, "simple")
self.assertEqual(cfg.model.history_encoder, "mlp")
self.assertEqual(cfg.train.epochs, 100)
def test_apply_overrides(self):
parser = build_argparser("test")
args = parser.parse_args(["--epochs", "50", "--lr", "0.01", "--backbone", "mobilenet"])
cfg = apply_overrides(Config(), args)
self.assertEqual(cfg.train.epochs, 50)
self.assertEqual(cfg.train.lr, 0.01)
self.assertEqual(cfg.model.cnn_backbone, "mobilenet")
def test_normalize_actions_flag(self):
"""--normalize-actions 和 --no-normalize-actions 都应可用"""
parser = build_argparser("test")
args = parser.parse_args(["--normalize-actions"])
self.assertTrue(args.normalize_actions)
parser = build_argparser("test")
args = parser.parse_args(["--no-normalize-actions"])
self.assertFalse(args.normalize_actions)
parser = build_argparser("test")
args = parser.parse_args([])
self.assertIsNone(args.normalize_actions)
if __name__ == "__main__":
unittest.main(verbosity=2)