File size: 2,032 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 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import torch
import numpy as np
import pytest
from ding.model import ConvEncoder, FCEncoder, IMPALAConvEncoder
from ding.torch_utils import is_differentiable
B = 4
C, H, W = 3, 128, 128
@pytest.mark.unittest
class TestEncoder:
def output_check(self, model, outputs):
loss = outputs.sum()
is_differentiable(loss, model)
def test_conv_encoder(self):
inputs = torch.randn(B, C, H, W)
model = ConvEncoder((C, H, W), hidden_size_list=[32, 48, 64, 64, 128], activation=torch.nn.Tanh())
print(model)
outputs = model(inputs)
self.output_check(model, outputs)
assert outputs.shape == (B, 128)
def test_dreamer_conv_encoder(self):
inputs = torch.randn(B, C, H, W)
model = ConvEncoder(
(C, H, W),
hidden_size_list=[32, 64, 128, 256, 128],
activation=torch.nn.SiLU(),
kernel_size=[4, 4, 4, 4],
layer_norm=True
)
print(model)
outputs = model(inputs)
self.output_check(model, outputs)
assert outputs.shape == (B, 128)
def test_fc_encoder(self):
inputs = torch.randn(B, 32)
hidden_size_list = [128 for _ in range(3)]
model = FCEncoder(32, hidden_size_list, res_block=True, activation=torch.nn.Tanh())
print(model)
outputs = model(inputs)
self.output_check(model, outputs)
assert outputs.shape == (B, hidden_size_list[-1])
hidden_size_list = [64, 128, 256]
model = FCEncoder(32, hidden_size_list, res_block=False, activation=torch.nn.Tanh())
print(model)
outputs = model(inputs)
self.output_check(model, outputs)
assert outputs.shape == (B, hidden_size_list[-1])
def test_impalaconv_encoder(self):
inputs = torch.randn(B, 3, 64, 64)
model = IMPALAConvEncoder(obs_shape=(3, 64, 64))
print(model)
outputs = model(inputs)
self.output_check(model, outputs)
assert outputs.shape == (B, 256)
|