Update.

[agtree2dot.git] / mlp.py

diff --git a/mlp.py b/mlp.py
index 963e7b7..4bf5841 100755 (executable)
--- a/mlp.py
+++ b/mlp.py
@@ -19,31 +19,32 @@
 
 import subprocess
 
+import torch
 from torch import nn
-from torch.nn import functional as fn
-from torch import Tensor
-from torch.autograd import Variable
 from torch.nn import Module
 
 import agtree2dot
 
 class MLP(Module):
     def __init__(self, input_dim, hidden_dim, output_dim):
-        super(MLP, self).__init__()
+        super().__init__()
         self.fc1 = nn.Linear(input_dim, hidden_dim)
         self.fc2 = nn.Linear(hidden_dim, output_dim)
 
     def forward(self, x):
         x = self.fc1(x)
-        x = fn.tanh(x)
+        x = torch.tanh(x)
         x = self.fc2(x)
         return x
 
 mlp = MLP(10, 20, 1)
-input = Variable(Tensor(100, 10).normal_())
-target = Variable(Tensor(100).normal_())
-output = mlp(input)
 criterion = nn.MSELoss()
+
+input = torch.randn(100, 10)
+target = torch.randn(100, 1)
+
+output = mlp(input)
+
 loss = criterion(output, target)
 
 agtree2dot.save_dot(loss,
@@ -61,15 +62,18 @@ agtree2dot.save_dot(loss,
 print('Generated mlp.dot')
 
 try:
-    fontname='Computer Modern'
-    fontsize=12
+
+    fontname = 'Computer Modern'
+    fontsize = 12
     subprocess.check_call(['dot', 'mlp.dot',
                            '-Lg',
                            '-T', 'pdf',
                            '-Efontname=' + fontname, '-Efontsize=' + str(fontsize),
                            '-Nfontname=' + fontname, '-Nfontsize=' + str(fontsize),
                            '-o', 'mlp.pdf' ])
+
 except subprocess.CalledProcessError:
+
     print('Calling the dot command failed. Is Graphviz installed?')
     sys.exit(1)