Add pass_state inference (#503)

* Add pass_state inference

* fix typo

* Improve coverage

CHANGELOG.md

@@ -18,9 +18,11 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Changed
 - Changed the default behaviour of the std metric to compute the sample std, in line with torch.std
 - Tqdm precision argument now rounds to decimal places rather than significant figures
+- Trial will now simply infer if the model has an argument called 'state'
 ### Deprecated
 ### Removed
 - Removed the old Model API (deprecated since version 0.2.0)
+- Removed the 'pass_state' argument from Trial, this will now be inferred
 ### Fixed
 - Fixed a bug in the weight decay callback which would result in potentially negative decay (now just uses torch.norm)
 - Fixed a bug in the cite decorator causing the citation to not show up correctly

tests/test_end_to_end.py

@@ -24,10 +24,17 @@ def f(self):
         out[2] = self.pars[2]-1
         return torch.sum(out**2)
 
-    def forward(self, _, state):
+    def forward(self, _):
         return self.f()
 
 
+class NetWithState(Net):
+    def forward(self, _, state=None):
+        if state is None:
+            raise ValueError
+        return super(NetWithState, self).forward(_)
+
+
 def loss(y_pred, y_true):
     return y_pred
 
@@ -37,10 +44,10 @@ def test_basic_opt(self):
         p = torch.tensor([2.0, 1.0, 10.0])
         training_steps = 1000
 
-        model = Net(p)
+        model = NetWithState(p)
         optim = torch.optim.SGD(model.parameters(), lr=0.01)
 
-        tbmodel = tb.Trial(model, optim, loss, pass_state=True).for_train_steps(training_steps)
+        tbmodel = tb.Trial(model, optim, loss).for_train_steps(training_steps).for_val_steps(1)
         tbmodel.run()
 
         self.assertAlmostEqual(model.pars[0].item(), 5.0, places=4)
@@ -54,17 +61,15 @@ def test_basic_checkpoint(self):
         model = Net(p)
         optim = torch.optim.SGD(model.parameters(), lr=0.01)
 
-        tbmodel = tb.Trial(model, optim, loss, callbacks=[tb.callbacks.MostRecent(filepath='test.pt')],
-                           pass_state=True).for_train_steps(training_steps)
+        tbmodel = tb.Trial(model, optim, loss, callbacks=[tb.callbacks.MostRecent(filepath='test.pt')]).for_train_steps(training_steps).for_val_steps(1)
         tbmodel.run(2)  # Simulate 2 'epochs'
 
         # Reload
         p = torch.tensor([2.0, 1.0, 10.0])
         model = Net(p)
         optim = torch.optim.SGD(model.parameters(), lr=0.01)
 
-        tbmodel = tb.Trial(model, optim, loss, callbacks=[tb.callbacks.MostRecent(filepath='test.pt')],
-                           pass_state=True).for_train_steps(training_steps)
+        tbmodel = tb.Trial(model, optim, loss, callbacks=[tb.callbacks.MostRecent(filepath='test.pt')]).for_train_steps(training_steps)
         tbmodel.load_state_dict(torch.load('test.pt'))
         self.assertEqual(len(tbmodel.state[tb.HISTORY]), 2)
         self.assertAlmostEqual(model.pars[0].item(), 5.0, places=4)