Make Model accept PreparedData rather than Data.

src/icosagon/model.py

@@ -1,8 +1,7 @@
 from .data import Data
 from typing import List, \
     Callable
-from .trainprep import prepare_training, \
-    TrainValTest
+from .trainprep import PreparedData
 import torch
 from .convlayer import DecagonLayer
 from .input import OneHotInputLayer
@@ -12,9 +11,8 @@ from .batch import PredictionsBatch
 
 
 class Model(object):
-    def __init__(self, data: Data,
+    def __init__(self, prep_d: PreparedData,
         layer_dimensions: List[int] = [32, 64],
-        ratios: TrainValTest = TrainValTest(.8, .1, .1),
         keep_prob: float = 1.,
         rel_activation: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
         layer_activation: Callable[[torch.Tensor], torch.Tensor] = torch.nn.functional.relu,
@@ -23,15 +21,12 @@ class Model(object):
         loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = torch.nn.functional.binary_cross_entropy_with_logits,
         batch_size: int = 100) -> None:
 
-        if not isinstance(data, Data):
-            raise TypeError('data must be an instance of Data')
+        if not isinstance(prep_d, PreparedData):
+            raise TypeError('prep_d must be an instance of PreparedData')
 
         if not isinstance(layer_dimensions, list):
             raise TypeError('layer_dimensions must be a list')
 
-        if not isinstance(ratios, TrainValTest):
-            raise TypeError('ratios must be an instance of TrainValTest')
-
         keep_prob = float(keep_prob)
 
         if not isinstance(rel_activation, FunctionType):
@@ -50,9 +45,8 @@ class Model(object):
 
         batch_size = int(batch_size)
 
-        self.data = data
+        self.prep_d = prep_d
         self.layer_dimensions = layer_dimensions
-        self.ratios = ratios
         self.keep_prob = keep_prob
         self.rel_activation = rel_activation
         self.layer_activation = layer_activation
@@ -61,15 +55,12 @@ class Model(object):
         self.loss = loss
         self.batch_size = batch_size
 
-        self.prep_d = None
         self.seq = None
         self.opt = None
 
         self.build()
 
     def build(self):
-        self.prep_d = prepare_training(self.data, self.ratios)
-
         in_layer = OneHotInputLayer(self.prep_d)
         last_output_dim = in_layer.output_dim
         seq = [ in_layer ]

tests/icosagon/test_model.py

@@ -5,7 +5,8 @@ from icosagon.trainprep import PreparedData, \
     PreparedRelationFamily, \
     PreparedRelationType, \
     TrainValTest, \
-    norm_adj_mat_one_node_type
+    norm_adj_mat_one_node_type, \
+    prepare_training
 import torch
 from icosagon.input import OneHotInputLayer
 from icosagon.convlayer import DecagonLayer
@@ -26,11 +27,12 @@ def test_model_01():
     fam = d.add_relation_family('Dummy-Dummy', 0, 0, False)
     fam.add_relation_type('Dummy Rel', torch.rand(10, 10).round())
 
-    m = Model(d)
+    prep_d = prepare_training(d, TrainValTest(.8, .1, .1))
 
-    assert m.data == d
+    m = Model(prep_d)
+
+    assert m.prep_d == prep_d
     assert m.layer_dimensions == [32, 64]
-    assert (m.ratios.train, m.ratios.val, m.ratios.test) == (.8, .1, .1)
     assert m.keep_prob == 1.
     assert _is_identity_function(m.rel_activation)
     assert m.layer_activation == torch.nn.functional.relu
@@ -38,7 +40,6 @@ def test_model_01():
     assert m.lr == 0.001
     assert m.loss == torch.nn.functional.binary_cross_entropy_with_logits
     assert m.batch_size == 100
-    assert isinstance(m.prep_d, PreparedData)
     assert isinstance(m.seq, torch.nn.Sequential)
     assert isinstance(m.opt, torch.optim.Optimizer)
 
@@ -50,7 +51,9 @@ def test_model_02():
     mat = torch.rand(10, 10).round().to_sparse()
     fam.add_relation_type('Dummy Rel', mat)
 
-    m = Model(d, ratios=TrainValTest(1., 0., 0.))
+    prep_d = prepare_training(d, TrainValTest(1., 0., 0.))
+
+    m = Model(prep_d)
 
     assert isinstance(m.prep_d, PreparedData)
     assert isinstance(m.prep_d.relation_families, list)
@@ -76,7 +79,9 @@ def test_model_03():
     mat = torch.rand(10, 10).round().to_sparse()
     fam.add_relation_type('Dummy Rel', mat)
 
-    m = Model(d, ratios=TrainValTest(1., 0., 0.))
+    prep_d = prepare_training(d, TrainValTest(1., 0., 0.))
+
+    m = Model(prep_d)
 
     state_dict = m.opt.state_dict()
     assert isinstance(state_dict, dict)
@@ -97,7 +102,9 @@ def test_model_04():
     fam.add_relation_type('Dummy Rel 1', mat)
     fam.add_relation_type('Dummy Rel 2', mat.clone())
 
-    m = Model(d, ratios=TrainValTest(1., 0., 0.))
+    prep_d = prepare_training(d, TrainValTest(1., 0., 0.))
+
+    m = Model(prep_d)
 
     assert len(list(m.seq[0].parameters())) == 1
     assert len(list(m.seq[1].parameters())) == 2
@@ -119,7 +126,9 @@ def test_model_05():
     fam.add_relation_type('Dummy Rel 2-1', mat)
     fam.add_relation_type('Dummy Rel 2-2', mat.clone())
 
-    m = Model(d, ratios=TrainValTest(1., 0., 0.))
+    prep_d = prepare_training(d, TrainValTest(1., 0., 0.))
+
+    m = Model(prep_d)
 
     assert len(list(m.seq[0].parameters())) == 1
     assert len(list(m.seq[1].parameters())) == 4
@@ -127,7 +136,7 @@ def test_model_05():
     assert len(list(m.seq[3].parameters())) == 6
 
 
-def test_model_05():
+def test_model_06():
     d = Data()
     d.add_node_type('Dummy', 10)
     d.add_node_type('Foobar', 20)
@@ -142,7 +151,9 @@ def test_model_05():
     fam.add_relation_type('Dummy Rel 2-1', mat)
     fam.add_relation_type('Dummy Rel 2-2', mat.clone())
 
-    m = Model(d, ratios=TrainValTest(1., 0., 0.))
+    prep_d = prepare_training(d, TrainValTest(1., 0., 0.))
+
+    m = Model(prep_d)
 
     assert len(list(m.seq[0].parameters())) == 2
     assert len(list(m.seq[1].parameters())) == 6
@@ -150,49 +161,53 @@ def test_model_05():
     assert len(list(m.seq[3].parameters())) == 6
 
 
-def test_model_06():
+def test_model_07():
     d = Data()
     d.add_node_type('Dummy', 10)
     fam = d.add_relation_family('Dummy-Dummy', 0, 0, False)
     fam.add_relation_type('Dummy Rel', torch.rand(10, 10).round())
 
+    prep_d = prepare_training(d, TrainValTest(.8, .1, .1))
+
     with pytest.raises(TypeError):
         m = Model(1)
 
     with pytest.raises(TypeError):
-        m = Model(d, layer_dimensions=1)
+        m = Model(prep_d, layer_dimensions=1)
 
     with pytest.raises(TypeError):
-        m = Model(d, ratios=1)
+        m = Model(prep_d, ratios=1)
 
     with pytest.raises(ValueError):
-        m = Model(d, keep_prob='x')
+        m = Model(prep_d, keep_prob='x')
 
     with pytest.raises(TypeError):
-        m = Model(d, rel_activation='x')
+        m = Model(prep_d, rel_activation='x')
 
     with pytest.raises(TypeError):
-        m = Model(d, layer_activation='x')
+        m = Model(prep_d, layer_activation='x')
 
     with pytest.raises(TypeError):
-        m = Model(d, dec_activation='x')
+        m = Model(prep_d, dec_activation='x')
 
     with pytest.raises(ValueError):
-        m = Model(d, lr='x')
+        m = Model(prep_d, lr='x')
 
     with pytest.raises(TypeError):
-        m = Model(d, loss=1)
+        m = Model(prep_d, loss=1)
 
     with pytest.raises(ValueError):
-        m = Model(d, batch_size='x')
+        m = Model(prep_d, batch_size='x')
 
 
-def test_model_07():
+def test_model_08():
     d = Data()
     d.add_node_type('Dummy', 10)
     fam = d.add_relation_family('Dummy-Dummy', 0, 0, False)
     fam.add_relation_type('Dummy Rel', torch.rand(10, 10).round())
 
-    m = Model(d)
+    prep_d = prepare_training(d, TrainValTest(.8, .1, .1))
+
+    m = Model(prep_d)
 
     m.run_epoch()