ntm: initial NTM implementation

Author: loudinthecloud

README.md

@@ -1,2 +1,48 @@
-# pytorch-ntm
-A Pytorch implementation of an NTM (Neural Turing Machine)
+# PyTorch Neural Turing Machine (NTM)
+
+PyTorch implementation of [Neural Turing Machines](https://arxiv.org/abs/1410.5401) (NTM).
+
+An **NTM** is a memory augumented neural network (attached to external memory) where the interactions with the external memory (address, read, write) are done using differentiable transformations. Overall, the network is end-to-end differentiable and thus trainable by a gradient based optimizer.
+
+The NTM is processing input in sequences, much like an LSTM, but with additional benfits: (1) The external memory allows the network to learn algorithmic tasks easier (2) Having a larger capacity without increasing the network's trainable parameters.
+
+The external memory allows the NTM to learn algorithmic tasks, that are much harder for LSTM to learn, and to maintain an internal state much longer than traditional LSTMs.
+
+## A PyTorch Implementation
+
+This repository implements a vanilla NTM in a straight forward way. The following architecture is used:
+
+![NTM Architecture](./images/ntm.png)
+
+* Support for batch leanring
+* Any read or write head configuration is supported (for example, 5 read heads and 2 write heads), the order of operation is specified by the user
+
+Example of training convergence for the **copy task** using 4 different seeds.
+
+![NTM Convergence](./images/train.png)
+
+ The following plot shows the cost per sequence length during training, the network was trained with `seed=7` and shows a fast convergence. Other seeds may not perform as well but should converge in less than 30K iterations.
+
+![NTM Convergence](./images/train2.png)
+
+Here is an animated GIF that shows how the model generalize. The model was evaluated after every 500 training samples, using the target sequence shown in the upper part of the image. The bottom part shows the network output at any given training stage.
+
+![NTM Convergence](./images/train-20.gif)
+
+The following is the same, but with `sequence length = 80`. Note that the network was trained with sequences of lengths 1 to 20.
+
+![NTM Convergence](./images/train-80.gif)
+
+
+## Installation
+
+The NTM can be used as a reusable module, currently not packaged though.
+
+1. Clone repository
+2. Install [PyTorch](http://pytorch.org/)
+3. pip install -r requirements.txt
+
+## Usage
+
+> python train.py
+

images/ntm.png

@@ -0,0 +1 @@
+__all__ = ['controller', 'head', 'memory', 'ntm']

images/train-20.gif

@@ -0,0 +1,47 @@
+"""LSTM Controller."""
+import torch
+from torch import nn
+from torch.nn import Parameter
+import numpy as np
+
+
+class LSTMController(nn.Module):
+    """An NTM controller based on LSTM."""
+    def __init__(self, num_inputs, num_outputs, num_layers):
+        super(LSTMController, self).__init__()
+
+        self.num_inputs = num_inputs
+        self.num_outputs = num_outputs
+        self.num_layers = num_layers
+
+        self.lstm = nn.LSTM(input_size=num_inputs,
+                            hidden_size=num_outputs,
+                            num_layers=num_layers)
+
+        # The hidden state is a learned parameter
+        self.lstm_h_bias = Parameter(torch.randn(self.num_layers, 1, self.num_outputs) * 0.05)
+        self.lstm_c_bias = Parameter(torch.randn(self.num_layers, 1, self.num_outputs) * 0.05)
+
+        self.reset_parameters()
+
+    def create_new_state(self, batch_size):
+        # Dimension: (num_layers * num_directions, batch, hidden_size)
+        lstm_h = self.lstm_h_bias.clone().repeat(1, batch_size, 1)
+        lstm_c = self.lstm_c_bias.clone().repeat(1, batch_size, 1)
+        return lstm_h, lstm_c
+
+    def reset_parameters(self):
+        for p in self.lstm.parameters():
+            if p.dim() == 1:
+                nn.init.constant(p, 0)
+            else:
+                stdev = 5 / (np.sqrt(self.num_inputs +  self.num_outputs))
+                nn.init.uniform(p, -stdev, stdev)
+
+    def size(self):
+        return self.num_inputs, self.num_outputs
+
+    def forward(self, x, prev_state):
+        x = x.unsqueeze(0)
+        outp, state = self.lstm(x, prev_state)
+        return outp.squeeze(0), state

images/train-80.gif

@@ -0,0 +1,130 @@
+"""NTM Read/Write Head."""
+import torch
+from torch import nn
+from torch.autograd import Variable
+import torch.nn.functional as F
+import numpy as np
+
+
+def _split_cols(mat, lengths):
+    """Split a 2D matrix to variable length columns."""
+    assert mat.size()[1] == sum(lengths), "Lengths must be summed to num columns"
+    l = np.cumsum([0] + lengths)
+    results = []
+    for s, e in zip(l[:-1], l[1:]):
+        results += [mat[:, s:e]]
+    return results
+
+
+class NTMHeadBase(nn.Module):
+    """An NTM Read/Write Head."""
+
+    def __init__(self, memory, controller_size):
+        """Initilize the read/write head.
+
+        :param memory: The :class:`NTMMemory` to be addressed by the head.
+        :param controller_size: The size of the internal representation.
+        """
+        super(NTMHeadBase, self).__init__()
+
+        self.memory = memory
+        self.N, self.M = memory.size()
+        self.controller_size = controller_size
+
+    def create_new_state(self, batch_size):
+        raise NotImplementedError
+
+    def init_weights(self):
+        raise NotImplementedError
+
+    def is_read_head(self):
+        return NotImplementedError
+
+    def _address_memory(self, k, β, g, s, γ, w_prev):
+        # Activations
+        k = F.relu(k)
+        β = F.relu(β)
+        g = F.sigmoid(g)
+        s = F.softmax(F.relu(s))
+        γ = 1 + F.relu(γ)
+
+        w = self.memory.address(k, β, g, s, γ, w_prev)
+
+        return w
+
+
+class NTMReadHead(NTMHeadBase):
+    def __init__(self, memory, controller_size):
+        super(NTMReadHead, self).__init__(memory, controller_size)
+
+        # Corresponding to k, β, g, s, γ sizes from the paper
+        self.read_lengths = [self.M, 1, 1, 3, 1]
+        self.fc_read = nn.Linear(controller_size, sum(self.read_lengths))
+        self.reset_parameters()
+
+    def create_new_state(self, batch_size):
+        # The state holds the previous time step address weightings
+        return Variable(torch.zeros(batch_size, self.N))
+
+    def reset_parameters(self):
+        # Initialize the linear layers
+        nn.init.xavier_uniform(self.fc_read.weight, gain=1.4)
+        nn.init.normal(self.fc_read.bias, std=0.01)
+
+    def is_read_head(self):
+        return True
+
+    def forward(self, embeddings, w_prev):
+        """NTMReadHead forward function.
+
+        :param embeddings: input representation of the controller.
+        :param w_prev: previous step state
+        """
+        o = self.fc_read(embeddings)
+        k, β, g, s, γ = _split_cols(o, self.read_lengths)
+
+        # Read from memory
+        w = self._address_memory(k, β, g, s, γ, w_prev)
+        r = self.memory.read(w)
+
+        return r, w
+
+
+class NTMWriteHead(NTMHeadBase):
+    def __init__(self, memory, controller_size):
+        super(NTMWriteHead, self).__init__(memory, controller_size)
+
+        # Corresponding to k, β, g, s, γ, e, a sizes from the paper
+        self.write_lengths = [self.M, 1, 1, 3, 1, self.M, self.M]
+        self.fc_write = nn.Linear(controller_size, sum(self.write_lengths))
+        self.reset_parameters()
+
+    def create_new_state(self, batch_size):
+        return Variable(torch.zeros(batch_size, self.N))
+
+    def reset_parameters(self):
+        # Initialize the linear layers
+        nn.init.xavier_uniform(self.fc_write.weight, gain=1.4)
+        nn.init.normal(self.fc_write.bias, std=0.01)
+
+    def is_read_head(self):
+        return False
+
+    def forward(self, embeddings, w_prev):
+        """NTMWriteHead forward function.
+
+        :param embeddings: input representation of the controller.
+        :param w_prev: previous step state
+        """
+        o = self.fc_write(embeddings)
+        k, β, g, s, γ, e, a = _split_cols(o, self.write_lengths)
+
+        # Handle activations
+        e = F.relu(e)
+        a = F.relu(a)
+
+        # Write to memory
+        w = self._address_memory(k, β, g, s, γ, w_prev)
+        self.memory.write(w, e, a)
+
+        return w

images/train.png

@@ -0,0 +1,102 @@
+"""An NTM's memory implementation."""
+import torch
+from torch.autograd import Variable
+import torch.nn.functional as F
+from torch import nn
+import numpy as np
+
+
+def _convolve(w, s):
+    """Circular convolution implementation."""
+    assert s.size(0) == 3
+    t = torch.cat([w[-2:], w, w[:2]])
+    c = F.conv1d(t.view(1, 1, -1), s.view(1, 1, -1)).view(-1)
+    return c[1:-1]
+
+
+class NTMMemory(nn.Module):
+    """Memory bank for NTM."""
+    def __init__(self, N, M):
+        """Initialize the NTM Memory matrix.
+
+        The memory's dimensions are (batch_size x N x M).
+        Each batch has it's own memory matrix.
+
+        :param N: Number of rows in the memory.
+        :param M: Number of columns/features in the memory.
+        """
+        super(NTMMemory, self).__init__()
+
+        self.N = N
+        self.M = M
+
+        # The memory bias allows the heads to learn how to initially address
+        # memory locations by content
+        self.mem_bias = Variable(torch.Tensor(N, M))
+        self.register_buffer('mem_bias', self.mem_bias.data)
+
+        # Initialize memory bias
+        stdev = 1 / (np.sqrt(N + M))
+        nn.init.uniform(self.mem_bias, -stdev, stdev)
+
+    def reset(self, batch_size):
+        """Initialize memory from bias, for start-of-sequence."""
+        self.batch_size = batch_size
+        self.memory = self.mem_bias.clone().repeat(batch_size, 1, 1)
+
+    def size(self):
+        return self.N, self.M
+
+    def read(self, w):
+        """Read from memory (according to section 3.1)."""
+        return torch.matmul(w.unsqueeze(1), self.memory).squeeze(1)
+
+    def write(self, w, e, a):
+        """write to memory (according to section 3.2)."""
+        self.prev_mem = self.memory
+        self.memory = Variable(torch.Tensor(self.batch_size, self.N, self.M))
+        for b in range(self.batch_size):
+            erase = torch.ger(w[b], e[b])
+            add = torch.ger(w[b], a[b])
+            self.memory[b] = self.prev_mem[b] * (1 - erase) + add
+
+    def address(self, k, β, g, s, γ, w_prev):
+        """NTM Addressing (according to section 3.3).
+
+        Returns a softmax weighting over the rows of the memory matrix.
+
+        :param k: The key vector.
+        :param β: The key strength (focus).
+        :param g: Scalar interpolation gate (with previous weighting).
+        :param s: Shift weighting.
+        :param γ: Sharpen weighting scalar.
+        :param w_prev: The weighting produced in the previous time step.
+        """
+        # Content focus
+        wc = self._similarity(k, β)
+
+        # Location focus
+        wg = self._interpolate(w_prev, wc, g)
+        ŵ = self._shift(wg, s)
+        w = self._sharpen(ŵ, γ)
+
+        return w
+
+    def _similarity(self, k, β):
+        k = k.view(self.batch_size, 1, -1)
+        w = F.softmax(β * F.cosine_similarity(self.memory + 1e-16, k + 1e-16, dim=-1))
+        return w
+
+    def _interpolate(self, w_prev, wc, g):
+        return g * wc + (1 - g) * w_prev
+
+    def _shift(self, wg, s):
+        result = Variable(torch.zeros(wg.size()))
+        for b in range(self.batch_size):
+            result[b] = _convolve(wg[b], s[b])
+        return result
+
+    def _sharpen(self, ŵ, γ):
+        w = ŵ ** γ
+        w = torch.div(w, torch.sum(w, dim=1).view(-1, 1) + 1e-16)
+        return w

images/train2.png

@@ -0,0 +1,95 @@
+#!/usr/bin/env python
+import torch
+from torch import nn
+from torch.autograd import Variable
+import torch.nn.functional as F
+
+
+class NTM(nn.Module):
+    """A Neural Turing Machine."""
+    def __init__(self, num_inputs, num_outputs, controller, memory, heads):
+        """Initialize the NTM.
+
+        :param num_inputs: External input size.
+        :param num_outputs: External output size.
+        :param controller: :class:`LSTMController`
+        :param memory: :class:`NTMMemory`
+        :param heads: list of :class:`NTMReadHead` or :class:`NTMWriteHead`
+
+        Note: This design allows the flexibility of using any number of read and
+              write heads independently, also, the order by which the heads are
+              called in controlled by the user (order in list)
+        """
+        super(NTM, self).__init__()
+
+        # Save arguments
+        self.num_inputs = num_inputs
+        self.num_outputs = num_outputs
+        self.controller = controller
+        self.memory = memory
+        self.heads = heads
+
+        self.N, self.M = memory.size()
+        _, self.controller_size = controller.size()
+
+        # Initialize the initial previous read values to random biases
+        self.num_read_heads = 0
+        self.init_r = []
+        for head in heads:
+            if head.is_read_head():
+                init_r_bias = Variable(torch.randn(1, self.M) * 0.01)
+                self.register_buffer("read{}_bias".format(self.num_read_heads), init_r_bias.data)
+                self.init_r += [init_r_bias]
+                self.num_read_heads += 1
+
+        assert self.num_read_heads > 0, "heads list must contain at least a single read head"
+
+        # Initialize a fully connected layer to produce the actual output:
+        #   [controller_output; previous_reads ] -> output
+        self.fc = nn.Linear(self.controller_size + self.num_read_heads * self.M, num_outputs)
+        self.reset_parameters()
+
+    def create_new_state(self, batch_size):
+        init_r = [r.clone().repeat(batch_size, 1) for r in self.init_r]
+        controller_state = self.controller.create_new_state(batch_size)
+        heads_state = [head.create_new_state(batch_size) for head in self.heads]
+
+        return init_r, controller_state, heads_state
+
+    def reset_parameters(self):
+        # Initialize the linear layer
+        nn.init.xavier_uniform(self.fc.weight, gain=1)
+        nn.init.normal(self.fc.bias, std=0.01)
+
+    def forward(self, x, prev_state):
+        """NTM forward function.
+
+        :param x: input vector (batch_size x num_inputs)
+        :param prev_state: The previous state of the NTM
+        """
+        # Unpack the previous state
+        prev_reads, prev_controller_state, prev_heads_states = prev_state
+
+        # Use the controller to get an embeddings
+        inp = torch.cat([x] + prev_reads, dim=1)
+        controller_outp, controller_state = self.controller(inp, prev_controller_state)
+
+        # Read/Write from the list of heads
+        reads = []
+        heads_states = []
+        for head, prev_head_state in zip(self.heads, prev_heads_states):
+            if head.is_read_head():
+                r, head_state = head(controller_outp, prev_head_state)
+                reads += [r]
+            else:
+                head_state = head(controller_outp, prev_head_state)
+            heads_states += [head_state]
+
+        # Generate Output
+        inp2 = torch.cat([controller_outp] + reads, dim=1)
+        o = F.sigmoid(self.fc(inp2))
+
+        # Pack the current state
+        state = (reads, controller_state, heads_states)
+
+        return o, state

ntm/__init__.py

@@ -0,0 +1,5 @@
+argcomplete
+attrs
+numpy
+pytest
+

ntm/controller.py

@@ -0,0 +1,240 @@
+"""Copy Task NTM model."""
+from attr import attrs, attrib, Factory
+import torch
+from torch import nn
+from torch.autograd import Variable
+from torch import optim
+import random
+import numpy as np
+
+from ntm.controller import LSTMController
+from ntm.memory import NTMMemory
+from ntm.head import NTMReadHead, NTMWriteHead
+from ntm.ntm import NTM
+
+
+# Encapsulation of the various NTM components for the Copy task
+class CopyTaskNTM(nn.Module):
+
+    def __init__(self, num_inputs, num_outputs, controller_size, controller_layers, N, M):
+        """Initialize an CopyTaskNTM.
+
+        :param num_inputs: External number of inputs.
+        :param num_outputs: External number of outputs.
+        :param controller_size: The size of the internal representation.
+        :param controller_layers: Number of controller layers.
+        :param N: Number of rows in the memory bank.
+        :param M: Number of cols/features in the memory bank.
+        """
+        super(CopyTaskNTM, self).__init__()
+
+        # Save args
+        self.num_inputs = num_inputs
+        self.num_outputs = num_outputs
+        self.controller_size = controller_size
+        self.N = N
+        self.M = M
+
+        # Create the NTM components
+        memory = NTMMemory(N, M)
+        heads = nn.ModuleList([
+            NTMReadHead(memory, controller_size),
+            NTMWriteHead(memory, controller_size)
+        ])
+        controller = LSTMController(num_inputs + M, controller_size, controller_layers)
+        self.ntm = NTM(num_inputs, num_outputs, controller, memory, heads)
+        self.memory = memory
+
+    def init_sequence(self, batch_size):
+        """Initializing the state."""
+        self.batch_size = batch_size
+        self.memory.reset(batch_size)
+        self.previous_state = self.ntm.create_new_state(batch_size)
+
+    def forward(self, x=None):
+        if x is None:
+            x = Variable(torch.zeros(self.batch_size, self.num_inputs))
+
+        o, self.previous_state = self.ntm(x, self.previous_state)
+        return o, self.previous_state
+
+    def calculate_num_params(self):
+        """Returns the total number of parameters."""
+        num_params = 0
+        for p in self.parameters():
+            num_params += p.data.view(-1).size(0)
+        return num_params
+
+
+# Generator of randomized test sequences
+def dataloader(num_batches,
+               batch_size,
+               seq_width,
+               min_len,
+               max_len):
+    """Generator of random sequences for the copy task.
+
+    Creates random batches of "bits" sequences.
+
+    All the sequences within each batch have the same length.
+    The length is [`min_len`, `max_len`]
+
+    :param num_batches: Total number of batches to generate.
+    :param seq_width: The width of each item in the sequence.
+    :param batch_size: Batch size.
+    :param min_len: Sequence minimum length.
+    :param max_len: Sequence maximum length.
+
+    NOTE: The input width is `seq_width + 1`, the additional input
+    contain the delimiter.
+    """
+    for batch_num in range(num_batches):
+
+        # All batches have the same sequence length
+        seq_len = random.randint(min_len, max_len)
+        seq = np.random.binomial(1, 0.5, (seq_len, batch_size, seq_width))
+        seq = Variable(torch.from_numpy(seq))
+
+        # The input includes an additional channel used for the delimiter
+        inp = Variable(torch.zeros(seq_len + 1, batch_size, seq_width + 1))
+        inp[:seq_len, :, :seq_width] = seq
+        inp[seq_len, :, seq_width] = 1.0 # delimiter in our control channel
+        outp = seq.clone()
+
+        yield batch_num+1, inp.float(), outp.float()
+
+
+def train_batch(net, criterion, optimizer, X, Y):
+    """Trains a single batch."""
+    optimizer.zero_grad()
+    seq_len, batch_size, _ = Y.size()
+
+    # New sequence
+    net.init_sequence(batch_size)
+
+    # Feed the sequence + delimiter
+    for i in range(seq_len+1):
+        net(X[i])
+
+    # Read the output (no input given)
+    y_out = Variable(torch.zeros(Y.size()))
+    for i in range(seq_len):
+        y_out[i], _ = net()
+
+    loss = criterion(y_out, Y)
+    loss.backward()
+    optimizer.step()
+
+    y_out_binarized = y_out.clone().data
+    y_out_binarized.apply_(lambda x: 0 if x < 0.5 else 1)
+
+    # The cost is the number of error bits per sequence
+    cost = torch.sum(torch.abs(y_out_binarized - Y.data))
+
+    return loss.data[0], cost / batch_size
+
+
+def evaluate(net, criterion, X, Y):
+    """Evaluate a single batch (without training)."""
+    seq_len, batch_size, _ = Y.size()
+
+    # New sequence
+    net.init_sequence(batch_size)
+
+    # Feed the sequence + delimiter
+    for i in range(seq_len+1):
+        o, _ = net(X[i])
+
+    # Read the output (no input given)
+    states = []
+    y_out = Variable(torch.zeros(Y.size()))
+    for i in range(seq_len):
+        y_out[i], state = net()
+        states += [state]
+
+    loss = criterion(y_out, Y)
+
+    y_out_binarized = y_out.clone().data
+    y_out_binarized.apply_(lambda x: 0 if x < 0.5 else 1)
+
+    # The cost is the number of error bits per sequence
+    cost = torch.sum(torch.abs(y_out_binarized - Y.data))
+
+    result = {
+        'loss': loss.data[0],
+        'cost': cost / batch_size,
+        'y_out': y_out,
+        'y_out_binarized': y_out_binarized,
+        'states': states
+    }
+
+    return result
+
+
+@attrs
+class CopyTaskParams(object):
+    name = attrib(default="copy-task")
+    controller_size = attrib(default=100)
+    controller_layers = attrib(default=1)
+    sequence_width = attrib(default=8)
+    sequence_min_len = attrib(default=1)
+    sequence_max_len = attrib(default=20)
+    memory_n = attrib(default=128)
+    memory_m = attrib(default=20)
+    num_batches = attrib(default=50000)
+    batch_size = attrib(default=1)
+    rmsprop_lr = attrib(default=1e-4)
+    rmsprop_momentum = attrib(default=0.9)
+    rmsprop_alpha = attrib(default=0.95)
+
+
+#
+# To create a network simply instantiate the `:class:CopyTaskModelTraining`,
+# all the components will be wired with the default values.
+# In case you'd like to change any of defaults, do the following:
+#
+# > params = CopyTaskParams(batch_size=4)
+# > model = CopyTaskModelTraining(params=params)
+#
+# Then use `model.net`, `model.optimizer` and `model.criterion` to train the
+# network. Call `model.train_batch` for training and `model.evaluate`
+# for evaluating.
+#
+# You may skip this alltogether, and use `:class:CopyTaskNTM` directly.
+#
+
+@attrs
+class CopyTaskModelTraining(object):
+    params = attrib(default=Factory(CopyTaskParams))
+    train_batch = attrib(default=train_batch)
+    evaluate = attrib(default=evaluate)
+    net = attrib()
+    dataloader = attrib()
+    criterion = attrib()
+    optimizer = attrib()
+
+    @net.default
+    def default_net(self):
+        # We have 1 additional input for the delimiter which is passed on a
+        # separate "control" channel
+        net = CopyTaskNTM(self.params.sequence_width + 1, self.params.sequence_width,
+                          self.params.controller_size, self.params.controller_layers,
+                          self.params.memory_n, self.params.memory_m)
+        return net
+
+    @dataloader.default
+    def default_dataloader(self):
+        return dataloader(self.params.num_batches, self.params.batch_size,
+                          self.params.sequence_width,
+                          self.params.sequence_min_len, self.params.sequence_max_len)
+
+    @criterion.default
+    def default_criterion(self):
+        return nn.BCELoss()
+
+    @optimizer.default
+    def default_optimizer(self):
+        return optim.RMSprop(self.net.parameters(),
+                             momentum=self.params.rmsprop_momentum,
+                             alpha=self.params.rmsprop_alpha,
+                             lr=self.params.rmsprop_lr)

ntm/head.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python
+# PYTHON_ARGCOMPLETE_OK
+"""Training for the Copy Task in Neural Turing Machines."""
+
+import argparse
+import json
+import time
+import random
+
+import argcomplete
+import torch
+import numpy as np
+
+from tasks.copytask import CopyTaskModelTraining
+
+
+# Default values for program arguments
+RANDOM_SEED = 1000
+REPORT_INTERVAL = 200
+CHECKPOINT_INTERVAL = 1000
+
+
+def get_ms():
+    """Returns the current time in miliseconds."""
+    return time.time() * 1000
+
+
+def init_seed(seed=None):
+    """Seed the RNGs for predicatability/reproduction purposes."""
+    if seed is None:
+        seed = int(get_ms() // 1000)
+
+    print("Using seed={}".format(seed))
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    random.seed(seed)
+
+
+def progress_clean():
+    """Clean the progress bar."""
+    print("\r{}".format(" " * 80), end='\r')
+
+
+def progress_bar(batch_num, report_interval, last_loss):
+    """Prints the progress until the next report."""
+    progress = (((batch_num-1) % report_interval) + 1) / report_interval
+    fill = int(progress * 40)
+    print("\r[{}{}]: {} (Loss: {:.4f})".format(
+        "=" * fill, " " * (40 - fill), batch_num, last_loss), end='')
+
+
+def save_checkpoint(net, name, seed, batch_num, losses, costs, seq_lengths):
+    progress_clean()
+
+    basename = "./{}-{}-batch-{}".format(name, seed, batch_num)
+    model_fname = basename + ".model"
+    print("Saving model checkpoint to: '{}'".format(model_fname))
+    torch.save(net.state_dict(), model_fname)
+
+    # Save the training history
+    train_fname = basename + ".json"
+    print("Saving model training history to '{}'".format(train_fname))
+    content = {
+        "loss": losses,
+        "cost": costs,
+        "seq_lengths": seq_lengths
+    }
+    open(train_fname, 'wt').write(json.dumps(content))
+
+
+def train_model(model,
+                args):
+
+    num_batches = model.params.num_batches
+    batch_size = model.params.batch_size
+
+    print("Training model for {} batches (batch_size={})...".format(
+        num_batches, batch_size))
+
+    losses = []
+    costs = []
+    seq_lengths = []
+    start_ms = get_ms()
+
+    for batch_num, x, y in model.dataloader:
+        loss, cost = model.train_batch(model.net, model.criterion, model.optimizer, x, y)
+        losses += [loss]
+        costs += [cost]
+        seq_lengths += [y.size(0)]
+
+        # Update the progress bar
+        progress_bar(batch_num, args.report_interval, loss)
+
+        # Report
+        if batch_num % args.report_interval == 0:
+            mean_loss = np.array(losses[-args.report_interval:]).mean()
+            mean_cost = np.array(costs[-args.report_interval:]).mean()
+            mean_time = int(((get_ms() - start_ms) / args.report_interval) / batch_size)
+            progress_clean()
+            print("Batch {} Loss: {:.6f} Cost: {:.2f} Time: {} ms/sequence".format(
+                batch_num, mean_loss, mean_cost, mean_time))
+            start_ms = get_ms()
+
+        # Checkpoint
+        if (args.checkpoint_interval != 0) and (batch_num % args.checkpoint_interval == 0):
+            save_checkpoint(model.net, model.params.name, args.seed,
+                            batch_num, losses, costs, seq_lengths)
+
+    print("Done training.")
+
+
+def init_arguments():
+    parser = argparse.ArgumentParser(prog='train.py')
+    parser.add_argument('--seed', type=int, default=RANDOM_SEED, help="Seed value for RNGs")
+    parser.add_argument('--checkpoint_interval', type=int, default=CHECKPOINT_INTERVAL,
+                        help="Checkpoint interval (in batches). 0 - disable")
+    parser.add_argument('--report_interval', type=int, default=REPORT_INTERVAL,
+                        help="Report interval (in batches)")
+
+    argcomplete.autocomplete(parser)
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    # Initialize arguments
+    args = init_arguments()
+
+    # Initialize random
+    init_seed(args.seed)
+
+    # Initialize the model
+    model = CopyTaskModelTraining()
+
+    print("Total number of parameters: {}".format(model.net.calculate_num_params()))
+    train_model(model, args)
+
+
+if __name__ == '__main__':
+    main()