1.First of all

It's very new, but I'll show you how to use the Neural ODE implementation library. By the way, Neural ODE is the best paper of NeuroIPS 2018.

This Neural ODE, authors have published a library of official repositories called torchdiffeq.

Although Neural ODE has many articles explaining the theory and interpretation, I thought that there were few Japanese articles describing the actual usage of this library, so I summarized the basic usage in this article. By the way, torchdiffeq is a library for PyTorch.

By reading this article, you will be able to:

--Torchdiffeq can solve the initial value problem of first-order ordinary differential equations --Torch diffeq can solve the initial value problem of second-order ordinary differential equations --The ODE Block layer, which is the layer that makes up Neural ODE, can be implemented.

2. Prerequisite knowledge

I will review the prerequisite knowledge for using torchdiffeq.

2.1 What is an ordinary differential equation?

Of the differential equations, the one with essentially only one unknown variable is called an ordinary differential equation. For example, differential equations such as $ \ frac {dz} {dt} = f (z (t), t) $ and $ m \ ddot {x} = -kx $ have multiple variables, but $ Since z $ and $ x $ are functions of $ t $, there is essentially only one unknown variable, $ t $, and it can be said that it is an ordinary differential equation.

2.2 What is Neural ODE?

There are many other easy-to-understand articles about Neural ODE, so please refer to them.

To briefly explain the outline, Neural ODE can be said to be a "neural network with continuous layers".

There is a theory that it is difficult to grasp the concept because there is a word called "ordinary differential equation" that is not heard in the neural network area, but I think the important point is that "layers are continuous". .. This makes it possible to, for example, "take out the output of the 0.5th layer", which was not possible with the conventional model.

Reference link below: -Neural network capable of expressing continuous dynamics -[NIPS 2018 Grand Prize Paper] From the University of Toronto: A completely new Neural Network model that connects the middle layers in a differentiable continuous space

3. How to use torchdiffeq

Now let's see how to use torchdiffeq.

3.1 Installation

To install, execute the following command.

pip install torchdiffeq

3.2 Example: First-order differential equation

Before actually implementing Neural ODE, let's take a simple first-order ODE as an example to see how to use torchdiffeq easily.

Consider the following differential equation. $ z(0) = 0, \\\ \frac{dz(t)}{dt} = f(t) = t $

This solution uses $ C $ as the constant of integration

\int dz = \int tdt+C　\\\ z(t) = \frac{t^2}{2} + C

Since $ z (0) = 0 , we can see that the solution of this differential equation is as follows. $ z(t) = \frac{t^2}{2} $$

Implementation by torchdiffeq

The simplest implementation to solve this problem with torchdiffeq is below.

`first_order.py`


from torchdiffeq import odeint

def func(t, z):
    return t

z0 = torch.Tensor([0])
t = torch.linspace(0,2,100)
out = odeint(func, z0, t)

Below, the points are itemized.

--The function func corresponds to $ f $ of the above differential equation $ \ frac {dz} {dt} = f (t, z) $. The arguments are in the order (t, z). The dimensions of the output must match z. You do not have to use either z or t. --The variable z0 is the initial value of dynamics. --The variable t represents time and must be a one-dimensional tensor, such astensor ([0., 0.1, 0.2, ..., 0.9, 1.0]). t [0] is the time corresponding to the initial value.

――It should be noted that the element of t must be a column that increases (decreases) monotonously in a narrow sense. An error will occur even if the same value is included, such as t = tensor ([0, 0, 1]). --Solve the differential equation with ʻodeint (func, z0, t) . The arguments are the function func, the initial value y0, and the time tin order. --The solver returns the value ofz at the time specified by t. That is, when t = tensor ([t0, t1, ..., tn]), ʻout = tensor ([z0, z1, ..., zn]). Since t [0] is the initial time, the output ʻout [0]always matchesz0`.

Plot the above results.

from matplotlib.pyplot as plt

plt.plot(t, out)
plt.axes().set_aspect('equal', 'datalim')  #Aspect ratio 1:Set to 1
plt.grid()
plt.xlim(0,2)
plt.show()

You can see that it matches the solution of the differential equation obtained by hand, $ z = \ frac {t ^ 2} {2} $.

3.3 (Reference) Example 2: Solving the second-order differential equation

Using torchdiffeq also solves second-order differential equations. As an example, we solve the (?) Simple vibration differential equation, which is familiar to science, with torchdiffeq. The differential equation of simple vibration is as follows. $ m\ddot{x} = -kx $ In the initial state, when $ t = 0 $, $ x = 1 $, $ \ dot {x} = \ frac {dx} {dt} = 0 $. The trick to solving the second-order differential equation is to decompose the second-order differential equation into two first-order differential equations. Specifically, do as follows.

\left[ \begin{array}{c} \dot{x} \\\ \ddot{x} \\\ \end{array} \right] = \left[ \begin{array}{cc} 0 & 1\\\ -\frac{k}{m} & 0\\\ \end{array} \right] \left[ \begin{array}{c} x \\\ \dot{x} \\\ \end{array} \right]

Here, $ \ boldsymbol {y} = \ left [ \begin{array}{c} x \
\dot{x} \
\end{array} If you put \ right] $, this second-order differential equation will result in the following first-order differential equation.

\frac{d\boldsymbol{y}}{dt} = f(\boldsymbol{y})

The implementation is as follows. $ k = 1, m = 1 $.

`oscillation.py`


class Oscillation:
    def __init__(self, km):
        self.mat = torch.Tensor([[0, 1],
                                 [-km, 0]])

    def solve(self, t, x0, dx0):
        y0 = torch.cat([x0, dx0])
        out = odeint(self.func, y0, t)
        return out

    def func(self, t, y):
        # print(t)
        out = y @ self.mat  # @Is a matrix product
        return out

if __name__=="__main__":
    x0 = torch.Tensor([1])
    dx0 = torch.Tensor([0])

    import numpy as np
    t = torch.linspace(0, 4 * np.pi, 1000)
    solver = Oscillation(1)
    out = solver.solve(t, x0, dx0)

When you draw it, you can see that the solution of simple vibration is properly obtained.

4. Implementation of ODE Block

Now that you are familiar with how to use torchdiffeq, let's see how to actually implement ODE Block. The ODE Block is a module that forms the dynamics of $ \ frac {dz} {dt} = f (t, z) $. The actual Neural ODE is constructed using ODE Block together with the normal Full-Connect layer and convolution layer.

The following implementation emphasizes simplicity and is just an example.

from torchdiffeq import odeint_adjoint as odeint

class ODEfunc(nn.Module):
    def __init__(self, dim):
        super(ODEfunc, self).__init__()
        self.seq = nn.Sequential(nn.Linear(dim, 124),
                                 nn.ReLU(),
                                 nn.Linear(124, 124),
                                 nn.ReLU(),
                                 nn.Linear(124, dim),
                                 nn.Tanh())

    def forward(self, t, x):
        out = self.seq(x)
        return out


class ODEBlock(nn.Module):
    def __init__(self, odefunc):
        super(ODEBlock, self).__init__()
        self.odefunc = odefunc
        self.integration_time = torch.tensor([0, 1]).float()

    def forward(self, x):
        self.integration_time = self.integration_time.type_as(x)
        out = odeint(self.odefunc, x, self.integration_time)
        return out[1]  # out[0]Because the initial value is included in.

To briefly explain,

--ODE Block treats the received input x as the initial value of the differential equation. ―― ʻODEfunc` is $ f $ that describes the dynamics of the system. --The integration interval of ODE Block is fixed at 0 ~ 1. And it returns the output of the layer at $ t = 1 $.

By doing this, you can use ODE Block as one module of the neural network as shown below.

class ODEnet(nn.Module):
    def __init__(self, in_dim, mid_dim, out_dim):
        super(ODEnet, self).__init__()

        odefunc = ODEfunc(dim=mid_dim)
        
        self.fc1 = nn.Linear(in_dim, mid_dim)
        self.relu1 = nn.ReLU(inplace=True)
        self.norm1 = nn.BatchNorm1d(mid_dim)
        self.ode_block = ODEBlock(odefunc)  #Use ODE Block
        self.norm2 = nn.BatchNorm1d(mid_dim)
        self.fc2 = nn.Linear(mid_dim, out_dim)

    def forward(self, x):
        batch_size = x.shape[0]
        x = x.view(batch_size, -1)

        out = self.fc1(x)
        out = self.relu1(out)
        out = self.norm1(out)
        out = self.ode_block(out)
        out = self.norm2(out)
        out = self.fc2(out)

        return out

This model was slow to calculate. However, using torchdiffeq doesn't always seem to slow it down, and as far as I've tried, the Neural ODE model in the official repository is as fast as a normal neural network. (This one should be a smaller model ...)

5. Summary

I introduced the rudimentary usage of torchdiffeq which is useful for implementing Neural ODE. If you would like to see the program that is actually training the model, please see the following Official torchdiffeq repository or [My implementation repository](https://github. com / TakadaTakumi / neuralODE_sample).

reference

torchdiffeq - GitHub My Implementation Repository

How to use the library "torchdiffeq" that implements Neural ODE's ODE Block