|
import pytest |
|
import torch |
|
|
|
from ding.torch_utils.distribution import Pd, CategoricalPd, CategoricalPdPytorch |
|
|
|
|
|
@pytest.mark.unittest |
|
class TestProbDistribution: |
|
|
|
def test_Pd(self): |
|
pd = Pd() |
|
with pytest.raises(NotImplementedError): |
|
pd.neglogp(torch.randn(5, )) |
|
with pytest.raises(NotImplementedError): |
|
pd.noise_mode() |
|
with pytest.raises(NotImplementedError): |
|
pd.mode() |
|
with pytest.raises(NotImplementedError): |
|
pd.sample() |
|
|
|
def test_CatePD(self): |
|
pd = CategoricalPd() |
|
logit1 = torch.randn(3, 5, requires_grad=True) |
|
logit2 = torch.randint(5, (3, ), dtype=torch.int64) |
|
|
|
pd.update_logits(logit1) |
|
entropy = pd.neglogp(logit2) |
|
assert entropy.requires_grad |
|
assert entropy.shape == torch.Size([]) |
|
|
|
entropy = pd.entropy() |
|
assert entropy.requires_grad |
|
assert entropy.shape == torch.Size([]) |
|
entropy = pd.entropy(reduction=None) |
|
assert entropy.requires_grad |
|
assert entropy.shape == torch.Size([3]) |
|
|
|
ret = pd.sample() |
|
assert ret.shape == torch.Size([3]) |
|
ret = pd.sample(viz=True) |
|
assert ret[0].shape == torch.Size([3]) |
|
|
|
ret = pd.mode() |
|
assert ret.shape == torch.Size([3]) |
|
ret = pd.mode(viz=True) |
|
assert ret[0].shape == torch.Size([3]) |
|
|
|
ret = pd.noise_mode() |
|
assert ret.shape == torch.Size([3]) |
|
ret = pd.noise_mode(viz=True) |
|
assert ret[0].shape == torch.Size([3]) |
|
|
|
pd = CategoricalPdPytorch() |
|
pd.update_logits(logit1) |
|
|
|
ret = pd.sample() |
|
assert ret.shape == torch.Size([3]) |
|
ret = pd.mode() |
|
assert ret.shape == torch.Size([3]) |
|
|
|
entropy = pd.entropy(reduction='mean') |
|
assert entropy.requires_grad |
|
assert entropy.shape == torch.Size([]) |
|
entropy = pd.entropy(reduction=None) |
|
assert entropy.requires_grad |
|
assert entropy.shape == torch.Size([3]) |
|
|
