adding load model

radionets-project · Sep 5, 2019 · 0f33437 · 0f33437
1 parent dacd248
commit 0f33437
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/mnist_cnn/mnist_cnn.py b/mnist_cnn/mnist_cnn.py
@@ -102,8 +102,14 @@ def get_model(data, lr=1e-1):
 run.fit(4000, learn)
 
 # Evaluate model
+from inspection import evaluate
 evaluate(valid_ds, learn.model)
 
 # Save model
 state = learn.model.state_dict()
-torch.save(state, './mnist_cnn_small_1')
+torch.save(state, './mnist_cnn_small_1.model')
+
+# +
+# Load model
+# m = learn.model
+# m.load_state_dict((torch.load('./mnist_cnn_small_1.model')))
diff --git a/mnist_cnn/mnist_cnn_big.py b/mnist_cnn/mnist_cnn_big.py
@@ -103,8 +103,18 @@ def get_model(data, lr=1e-1): #1e-1
 run.fit(4000, learn)
 
 # Evaluate model
+from inspection import evaluate
 evaluate(valid_ds, learn.model)
 
+# +
 # Save model
-state = learn.model.state_dict()
-torch.save(state, './mnist_cnn_small_1')
+# state = learn.model.state_dict()
+# torch.save(state, './mnist_cnn_big_1.model')
+
+# +
+# Load model
+# m = learn.model
+# m.load_state_dict((torch.load('./mnist_cnn_big_1.model')))
+# -
+
+