File size: 957 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
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])
|