-
Notifications
You must be signed in to change notification settings - Fork 0
/
trainer_test.py
111 lines (89 loc) · 3.71 KB
/
trainer_test.py
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
# -*- coding: utf-8 -*-
# Copyright (C) 2019 Machine Learning Group of the University of Oldenburg.
# Licensed under the Academic Free License version 3.0
import tvo
from tvo.trainer import Trainer
from tvo.models import NoisyOR
from tvo.variational import RandomSampledVarStates, FullEM
from tvo.utils.data import TVODataLoader
import pytest
import torch as to
import numpy as np
@pytest.fixture(
scope="function", params=[pytest.param(tvo.get_device().type, marks=pytest.mark.gpu)]
)
def setup(request):
class Setup:
device = tvo.get_device()
precision = to.float32
N, D, S, H, S_new = 10, 16, 8, 8, 10
model = NoisyOR(H, D, precision=precision)
_td = to.randint(2, size=(N, D), dtype=to.uint8, device=device)
data = TVODataLoader(_td, batch_size=N)
_td = to.randint(2, size=(N, D), dtype=to.uint8, device=device)
test_data = TVODataLoader(_td, batch_size=N)
var_states = RandomSampledVarStates(N, H, S, precision, S_new)
test_states = RandomSampledVarStates(N, H, S, precision, S_new)
return Setup
def test_training(setup):
trainer = Trainer(setup.model, setup.data, setup.var_states)
d1 = trainer.em_step()
d2 = trainer.em_step()
assert "train_F" in d1 and "train_subs" in d1
assert "test_F" not in d1 and "test_subs" not in d1
assert d1["train_F"] < d2["train_F"]
def test_training_with_tensor(setup):
data = to.randint(2, size=(setup.N, setup.D)).byte()
trainer = Trainer(setup.model, data, setup.var_states)
d1 = trainer.em_step()
d2 = trainer.em_step()
assert "train_F" in d1 and "train_subs" in d1
assert "test_F" not in d1 and "test_subs" not in d1
assert d1["train_F"] < d2["train_F"]
def test_training_with_valid(setup):
trainer = Trainer(
setup.model,
train_data=setup.data,
train_states=setup.var_states,
test_data=setup.test_data,
test_states=setup.test_states,
)
d1 = trainer.em_step()
d2 = trainer.em_step()
assert "train_F" in d1 and "train_subs" in d1
assert "test_F" in d1 and "test_subs" in d1
assert d1["train_F"] < d2["train_F"]
def test_testing(setup):
trainer = Trainer(setup.model, test_data=setup.test_data, test_states=setup.test_states)
d1 = trainer.em_step()
d2 = trainer.em_step()
assert "train_F" not in d1 and "train_subs" not in d1
assert "test_F" in d1 and "test_subs" in d1
assert d1["test_F"] < d2["test_F"]
def test_estep(setup):
trainer = Trainer(setup.model, test_data=setup.test_data, test_states=setup.test_states)
d1 = trainer.e_step()
d2 = trainer.e_step()
assert d1["test_F"] < d2["test_F"]
@pytest.mark.gpu
def test_rollback():
# A seed for which we know the unconstrained M-step will decrease free energies.
to.manual_seed(123)
eps = 1e-6
N, H, D = 10, 4, 32
data = (to.rand(N, D, device=tvo.get_device()) < 0.8).byte() # noisy data
dataloader = TVODataLoader(data, batch_size=N)
# make two identical copies of the model: we'll train twice with same initial conditions
# parameters are initialized stochastically.
m1 = NoisyOR(H, D, precision=to.float64)
m2 = NoisyOR(H, D, precision=m1.precision, W_init=m1.theta["W"], pi_init=m1.theta["pies"])
# use FullEM so training is deterministic
states = FullEM(N, H, m1.precision)
# without rollback, we get decreasing free energies
trainer = Trainer(m1, dataloader, states)
F = [trainer.em_step()["train_F"] for _ in range(50)]
assert np.any(np.diff(F) < -eps)
# with rollback, F always increases
trainer = Trainer(m2, dataloader, states, rollback_if_F_decreases=["W"])
F = [trainer.em_step()["train_F"] for _ in range(50)]
assert np.any(np.diff(F) >= -eps)