257 lines
9.8 KiB
Python
257 lines
9.8 KiB
Python
"""单元测试：pipeline.packing 模块中的 SequencePacker 类"""
|
||
|
||
import pytest
|
||
import torch
|
||
from pipeline.packing import SequencePacker
|
||
|
||
|
||
class TestSequencePacker:
|
||
"""SequencePacker 类的测试套件"""
|
||
|
||
def test_normal_packing(self):
|
||
"""测试正常打包场景：多个序列正确打包成固定长度的包"""
|
||
packer = SequencePacker(pack_size=10, pad_value=0)
|
||
|
||
sequences = [
|
||
torch.tensor([1, 2, 3], dtype=torch.int32),
|
||
torch.tensor([4, 5], dtype=torch.int32),
|
||
torch.tensor([6, 7, 8, 9], dtype=torch.int32),
|
||
]
|
||
|
||
packages = packer.pack(sequences)
|
||
|
||
# 验证至少有包输出
|
||
assert len(packages) >= 1
|
||
|
||
# 验证每个包的长度是正确的
|
||
for pkg in packages:
|
||
assert pkg.shape == (10,)
|
||
# 验证填充值
|
||
# 检查所有非零元素都在前几个位置，或者包是满的
|
||
non_zero_count = (pkg != 0).sum().item()
|
||
# 非零元素的数量应该等于原始序列元素的总和
|
||
total_elements = sum(s.numel() for s in sequences)
|
||
# 由于打包，第一个包包含3+2=5个元素，第二个包包含4个元素
|
||
# 第一个包应该包含前两个序列
|
||
pkg1 = packages[0]
|
||
# 序列[1,2,3]和[4,5]按长度降序排序后是[1,2,3]在前，然后[4,5]
|
||
# 但排序是原地修改...等等，我们已经修复了使用sorted()
|
||
# 所以排序后的顺序是[6,7,8,9], [1,2,3], [4,5]
|
||
# 第一个包包含[6,7,8,9]和部分[1,2,3] = 4+3=7，剩余3个位置放[4,5]
|
||
# 所以第一个包应该是[6,7,8,9,1,2,3,4,5,0]
|
||
|
||
# 简化测试：验证打包后的张量包含所有原始数据
|
||
all_values = []
|
||
for pkg in packages:
|
||
non_zero = pkg[pkg != 0].tolist()
|
||
all_values.extend(non_zero)
|
||
|
||
# 检查所有原始数据是否都被包含
|
||
original_values = [1, 2, 3, 4, 5, 6, 7, 8, 9]
|
||
for val in original_values:
|
||
assert val in all_values, f"Value {val} not found in packages"
|
||
|
||
def test_empty_list_input(self):
|
||
"""测试空列表输入"""
|
||
packer = SequencePacker(pack_size=10)
|
||
|
||
packages = packer.pack([])
|
||
|
||
assert packages == []
|
||
|
||
# 验证内部状态已正确初始化
|
||
assert packer._current_pack is not None
|
||
assert packer._current_pos == 0
|
||
|
||
def test_single_sequence_input(self):
|
||
"""测试单个序列输入"""
|
||
packer = SequencePacker(pack_size=10, pad_value=-1)
|
||
|
||
sequences = [torch.tensor([1, 2, 3], dtype=torch.int32)]
|
||
|
||
packages = packer.pack(sequences)
|
||
|
||
assert len(packages) == 1
|
||
pkg = packages[0]
|
||
assert pkg.shape == (10,)
|
||
assert pkg[:3].tolist() == [1, 2, 3]
|
||
assert pkg[3:].tolist() == [-1] * 7
|
||
|
||
def test_truncate_long_sequence(self, caplog):
|
||
"""测试超长序列截断，验证警告日志是否触发"""
|
||
packer = SequencePacker(pack_size=5, pad_value=0)
|
||
|
||
sequences = [
|
||
torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32), # 长度8，超过pack_size=5
|
||
]
|
||
|
||
packages = packer.pack(sequences)
|
||
|
||
assert len(packages) == 1
|
||
pkg = packages[0]
|
||
assert pkg.shape == (5,)
|
||
assert pkg.tolist() == [1, 2, 3, 4, 5] # 只保留前5个元素
|
||
|
||
# 验证警告日志已触发
|
||
assert "truncating" in caplog.text.lower() or "exceeds" in caplog.text.lower()
|
||
|
||
def test_padding_value(self):
|
||
"""测试填充值正确应用"""
|
||
packer = SequencePacker(pack_size=8, pad_value=99)
|
||
|
||
sequences = [
|
||
torch.tensor([1, 2], dtype=torch.int32),
|
||
torch.tensor([3], dtype=torch.int32),
|
||
]
|
||
|
||
packages = packer.pack(sequences)
|
||
|
||
assert len(packages) == 1
|
||
pkg = packages[0]
|
||
|
||
# 前3个元素是数据
|
||
assert pkg[:3].tolist() == [1, 2, 3]
|
||
# 后5个元素是填充值
|
||
assert pkg[3:].tolist() == [99] * 5
|
||
|
||
def test_different_dtypes(self):
|
||
"""测试支持不同 dtype (int32, int64, float32)"""
|
||
# int32
|
||
packer_int32 = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32)
|
||
sequences_int32 = [torch.tensor([1, 2, 3], dtype=torch.int32)]
|
||
packages_int32 = packer_int32.pack(sequences_int32)
|
||
assert packages_int32[0].dtype == torch.int32
|
||
|
||
# int64
|
||
packer_int64 = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int64)
|
||
sequences_int64 = [torch.tensor([1, 2, 3], dtype=torch.int64)]
|
||
packages_int64 = packer_int64.pack(sequences_int64)
|
||
assert packages_int64[0].dtype == torch.int64
|
||
|
||
# float32
|
||
packer_float32 = SequencePacker(pack_size=10, pad_value=0.0, dtype=torch.float32)
|
||
sequences_float32 = [torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)]
|
||
packages_float32 = packer_float32.pack(sequences_float32)
|
||
assert packages_float32[0].dtype == torch.float32
|
||
|
||
def test_non_1d_tensor_raises_error(self):
|
||
"""测试非1D张量是否抛出异常"""
|
||
packer = SequencePacker(pack_size=10)
|
||
|
||
# 2D 张量应该抛出异常
|
||
sequences_2d = [torch.tensor([[1, 2], [3, 4]])] # shape: (2, 2)
|
||
with pytest.raises(ValueError, match="Expected 1D tensor"):
|
||
packer.pack(sequences_2d)
|
||
|
||
# 0D 张量 (标量) 应该抛出异常
|
||
sequences_0d = [torch.tensor(5)] # shape: ()
|
||
with pytest.raises(ValueError, match="Expected 1D tensor"):
|
||
packer.pack(sequences_0d)
|
||
|
||
# 3D 张量应该抛出异常
|
||
sequences_3d = [torch.tensor([[[1, 2]]])] # shape: (1, 1, 2)
|
||
with pytest.raises(ValueError, match="Expected 1D tensor"):
|
||
packer.pack(sequences_3d)
|
||
|
||
def test_reset_method(self):
|
||
"""测试 reset() 方法是否正确重置内部状态"""
|
||
packer = SequencePacker(pack_size=10, pad_value=0)
|
||
|
||
# 第一次打包
|
||
sequences1 = [torch.tensor([1, 2, 3], dtype=torch.int32)]
|
||
packer.pack(sequences1)
|
||
|
||
# 验证内部状态已更新
|
||
assert packer._current_pos == 0
|
||
assert packer._current_pack is None # 最后一个包已发送，设置为None
|
||
|
||
# 重置
|
||
packer.reset()
|
||
|
||
# 验证重置后的状态
|
||
assert packer._current_pos == 0
|
||
assert packer._current_pack is not None
|
||
assert packer._current_pack.shape == (10,)
|
||
assert packer._current_pack.tolist() == [0] * 10
|
||
|
||
# 验证重置后可以继续正常使用
|
||
sequences2 = [torch.tensor([4, 5, 6], dtype=torch.int32)]
|
||
packages = packer.pack(sequences2)
|
||
|
||
assert len(packages) == 1
|
||
assert packages[0][:3].tolist() == [4, 5, 6]
|
||
|
||
def test_input_list_not_modified(self):
|
||
"""测试输入列表是否未被修改（使用 sorted 而非 sort）"""
|
||
packer = SequencePacker(pack_size=10)
|
||
|
||
# 创建原始序列列表（故意不按长度排序）
|
||
original_sequences = [
|
||
torch.tensor([3], dtype=torch.int32), # 长度1
|
||
torch.tensor([1, 2], dtype=torch.int32), # 长度2
|
||
torch.tensor([4, 5, 6, 7], dtype=torch.int32), # 长度4
|
||
]
|
||
|
||
# 保存原始顺序的字符串表示
|
||
original_repr = [seq.tolist() for seq in original_sequences]
|
||
|
||
# 打包
|
||
packer.pack(original_sequences)
|
||
|
||
# 验证输入列表未被修改
|
||
current_repr = [seq.tolist() for seq in original_sequences]
|
||
assert current_repr == original_repr, "输入列表被修改了，应该使用 sorted() 而非 sort()"
|
||
|
||
def test_exact_pack_size_fit(self):
|
||
"""测试序列长度恰好等于 pack_size 的情况"""
|
||
packer = SequencePacker(pack_size=5, pad_value=0)
|
||
|
||
sequences = [
|
||
torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32),
|
||
torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32),
|
||
]
|
||
|
||
packages = packer.pack(sequences)
|
||
|
||
# 每个序列恰好占满一个包
|
||
assert len(packages) == 2
|
||
|
||
assert packages[0].tolist() == [1, 2, 3, 4, 5]
|
||
assert packages[1].tolist() == [6, 7, 8, 9, 10]
|
||
|
||
def test_multiple_packs_full_utilization(self):
|
||
"""测试多个包的高效利用"""
|
||
packer = SequencePacker(pack_size=10, pad_value=-1)
|
||
|
||
# 创建多个小序列，确保高效打包
|
||
sequences = [
|
||
torch.tensor([1], dtype=torch.int32),
|
||
torch.tensor([2], dtype=torch.int32),
|
||
torch.tensor([3], dtype=torch.int32),
|
||
torch.tensor([4], dtype=torch.int32),
|
||
torch.tensor([5], dtype=torch.int32),
|
||
torch.tensor([6], dtype=torch.int32),
|
||
torch.tensor([7], dtype=torch.int32),
|
||
torch.tensor([8], dtype=torch.int32),
|
||
torch.tensor([9], dtype=torch.int32),
|
||
torch.tensor([10], dtype=torch.int32),
|
||
torch.tensor([11], dtype=torch.int32),
|
||
]
|
||
|
||
packages = packer.pack(sequences)
|
||
|
||
# 前10个序列打包成一个包，最后一个序列单独一个包
|
||
assert len(packages) == 2
|
||
assert packages[0].tolist() == list(range(1, 11))
|
||
assert packages[1].tolist() == [11] + [-1] * 9
|
||
|
||
def test_dtype_mismatch_warning(self, caplog):
|
||
"""测试 dtype 不匹配时的警告"""
|
||
packer = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32)
|
||
|
||
sequences = [torch.tensor([1, 2, 3], dtype=torch.int64)]
|
||
|
||
packages = packer.pack(sequences)
|
||
|
||
# 应该触发 dtype 不匹配警告
|
||
assert "dtype" in caplog.text.lower() or "converted" in caplog.text.lower() |