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)
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))