mlp.py

   1 #!/usr/bin/env python
   2
   3 #########################################################################
   4 # This program is free software: you can redistribute it and/or modify  #
   5 # it under the terms of the version 3 of the GNU General Public License #
   6 # as published by the Free Software Foundation.                         #
   7 #                                                                       #
   8 # This program is distributed in the hope that it will be useful, but   #
   9 # WITHOUT ANY WARRANTY; without even the implied warranty of            #
  10 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU      #
  11 # General Public License for more details.                              #
  12 #                                                                       #
  13 # You should have received a copy of the GNU General Public License     #
  14 # along with this program. If not, see <http://www.gnu.org/licenses/>.  #
  15 #                                                                       #
  16 # Written by and Copyright (C) Francois Fleuret                         #
  17 # Contact <francois.fleuret@idiap.ch> for comments & bug reports        #
  18 #########################################################################
  19
  20 from torch import nn
  21 from torch.nn import functional as fn
  22 from torch import Tensor
  23 from torch.autograd import Variable
  24 from torch.nn import Module
  25
  26 import agtree2dot
  27
  28 class MLP(Module):
  29     def __init__(self, input_dim, hidden_dim, output_dim):
  30         super(MLP, self).__init__()
  31         self.fc1 = nn.Linear(input_dim, hidden_dim)
  32         self.fc2 = nn.Linear(hidden_dim, output_dim)
  33
  34     def forward(self, x):
  35         x = self.fc1(x)
  36         x = fn.tanh(x)
  37         x = self.fc2(x)
  38         return x
  39
  40 mlp = MLP(10, 20, 1)
  41 input = Variable(Tensor(100, 10).normal_())
  42 target = Variable(Tensor(100).normal_())
  43 output = mlp(input)
  44 criterion = nn.MSELoss()
  45 loss = criterion(output, target)
  46
  47 agtree2dot.save_dot(loss,
  48                     { input: 'input', target: 'target', loss: 'loss' },
  49                     open('./mlp.dot', 'w'))
  50
  51 print('Generated mlp.dot. You can convert it to pdf with')
  52 print('> dot mlp.dot -Lg -T pdf -o mlp.pdf')