Code

To use torch-train into your own project, you can use it in place of the torch.nn.Module. Here we show some simple examples on how to use the torch-train Module in your own python code. For a complete documentation we refer to the Reference guide.

Import

To import the Module use

from torchtrain import Module

Working example

In this example, we create a basic torch Module and use its fit() and predict() methods to train and test. First we import torch and the torchtrain Module

# imports
import torch
import torch.nn as nn
from   torchtrain import Module

Next we create our simple network consisting of 2 layers and a softmax output function.

Note

We extend the torchtrain.Module instead of the torch.nn.Module like you normally would.

Furthermore we implement the forward() method to propagate through the network.

class MyNetwork(Module):

    def __init__(self, size_input, size_hidden, size_output):
        """Create simple network"""
        # Initialise super
        super().__init__()

        # Set layers
        self.layer_1 = nn.Linear(size_input , size_hidden)
        self.layer_2 = nn.Linear(size_hidden, size_output)
        self.softmax = nn.LogSoftmax(dim=1)

    def forward(X):
        """Forward through network"""
        # Propagate layer 1
        out = self.layer_1(X)
        # Propagate layer 2
        out = self.layer_2(out)
        # Propagate softmax layer
        out = self.softmax(out)
        # Return result
        return out

Now that we have created our network, we generate some random training and testing data.

# Generate random data
X_train = torch.rand((1024, 10))
y_train = (torch.rand(1024)*10).to(torch.int64)
X_test  = torch.rand((1024, 10))
y_test  = (torch.rand(1024)*10).to(torch.int64)

Finally, we create the network and invoke its fit() and predict() methods.

# Create network
net = MyNetwork(10, 128, 10)

# Fit network
net.fit(X_train, y_train,
    epochs        = 10,
    batch_size    = 32,
    learning_rate = 0.01,
    criterion     = nn.NLLLoss(),
    optimizer     = optim.SGD,
    variable      = False,
    verbose       = True
)

# Predict network
y_pred = net.predict(X_test,
    batch_size = 32,
    variable   = False,
    verbose    = True
)