|
import pytest |
|
import torch |
|
from ding.torch_utils import is_differentiable |
|
|
|
from lzero.model.common import RepresentationNetwork |
|
|
|
|
|
@pytest.mark.unittest |
|
class TestCommon: |
|
|
|
def output_check(self, model, outputs): |
|
if isinstance(outputs, torch.Tensor): |
|
loss = outputs.sum() |
|
elif isinstance(outputs, list): |
|
loss = sum([t.sum() for t in outputs]) |
|
elif isinstance(outputs, dict): |
|
loss = sum([v.sum() for v in outputs.values()]) |
|
is_differentiable(loss, model) |
|
|
|
@pytest.mark.parametrize('batch_size', [10]) |
|
def test_representation_network(self, batch_size): |
|
batch = batch_size |
|
obs = torch.rand(batch, 1, 3, 3) |
|
representation_network = RepresentationNetwork( |
|
observation_shape=[1, 3, 3], num_res_blocks=1, num_channels=16, downsample=False |
|
) |
|
state = representation_network(obs) |
|
assert state.shape == torch.Size([10, 16, 3, 3]) |
|
|
