Ejemplo n.º 1
0
    def test_soft_eval(self):
        args = default_args()
        args['test_temperature_multiplier'] = 0.0000001
        model = MockModel(Pyramid, args)

        X, transitions = get_batch()

        class Projection(nn.Module):
            def forward(self, x):
                return x[:, :default_args()['model_dim']]

        class mlp(nn.Module):
            def forward(self, x):
                return x

        def selection_fn(selection_input):
            # Compose in random order
            return Variable(torch.rand(selection_input.data.size()[0], 1))

        class Reduce(nn.Module):
            def forward(self, lefts, rights):
                return lefts + rights

        model.encode = Projection()
        model.composition_fn = Reduce()
        model.selection_fn = selection_fn
        model.mlp = mlp()

        model.eval()
        outputs = model(X, transitions)

        assert outputs[0][0].data[0] == (3 + 1 + 2 + 1)
        assert outputs[1][0].data[0] == (3 + 2 + 4 + 5)
Ejemplo n.º 2
0
    def test_basic_stack(self):
        model = MockModel(BaseModel, default_args())

        X, transitions = get_batch()

        class Projection(nn.Module):
            def forward(self, x):
                return x[:, :default_args()['model_dim']]

        class Reduce(nn.Module):
            def forward(self, lefts, rights, tracking):
                batch_size = len(lefts)
                return torch.chunk(
                    torch.cat(lefts, 0) - torch.cat(rights, 0), batch_size, 0)

        model.encode = Projection()
        model.spinn.reduce = Reduce()

        model(X, transitions)
        outputs = model.spinn_outp[0]

        assert outputs[0][0].data[0] == (3 - (1 - (2 - 1)))
        assert outputs[1][0].data[0] == ((3 - 2) - (4 - 5))