torchutils 0.0.4
PyTorch utility APIs.
License: MIT
Keywords: machine-learning deep-learning pytorch neuralnetwork

==========
TorchUtils
==========

.. image:: https://img.shields.io/pypi/v/torchutils?color=success
    :target: https://img.shields.io/pypi/v/torchutils
    :alt: PyPI

.. image:: https://travis-ci.org/anjandeepsahni/torchutils.svg?branch=master
    :target: https://travis-ci.org/anjandeepsahni/torchutils
    :alt: Build Status

.. image:: https://codecov.io/gh/anjandeepsahni/torchutils/branch/master/graph/badge.svg
    :target: https://codecov.io/gh/anjandeepsahni/torchutils
    :alt: Code Coverage

.. image:: https://img.shields.io/github/release-date/anjandeepsahni/torchutils?color=informational&label=release%20date
    :target: https://img.shields.io/github/release-date/anjandeepsahni/torchutils
    :alt: Release Date

.. image:: https://img.shields.io/github/license/anjandeepsahni/torchutils?color=informational
    :target: https://img.shields.io/github/license/anjandeepsahni/torchutils
    :alt: License

.. image:: https://pepy.tech/badge/torchutils
    :target: https://pepy.tech/badge/torchutils
    :alt: Downloads

|

**TorchUtils** is a Python package providing helpful utility APIs for your
PyTorch projects.

Features
--------

* Save/load checkpoints_.
* Calculate dataset statistics_ (mean, std, var). Also calculate and track running statistics of data.
* Get/set `learning rate`_.
* Track `evaluation metrics`_ such as accuracy, running loss, hamming loss.
* Print `model summary`_. Supports: Linear/MLP, Convolution Network, Recurrent Network (RNN/LSTM/GRU), Recursive Network.
* Calculate `model FLOPs`_.
* Calculate total `model parameters`_.
* Set `random seed`_.
* Visualize `gradient flow`_ in your network.

Requirements
------------

* PyTorch >= 1.0.0
* Numpy >= 1.16.2
* Matplotlib >= 3.0.3

Installation
------------

PyPi::

    $ pip install torchutils

Conda::

    $ conda install -c sahni torchutils

Documentation
-------------
Detailed API documentation is available here_.

.. _here: https://anjandeepsahni.github.io/torchutils/readme.html

Examples
--------

.. _checkpoints:

Checkpoint:
^^^^^^^^^^^

.. code:: python

    import torchvision
    import torchutils as tu
    import torch.optim as optim

    model = torchvision.models.alexnet()
    optimizer = optim.Adam(model.parameters())
    scheduler = optim.lr_scheduler.ExponentialLR(optimizer, 0.1)
    print('Original learning rate:', tu.get_lr(optimizer))

    # load checkpoint
    start_epoch = tu.load_checkpoint(model_path='.',
                            ckpt_name='model_20190814-212442_e0_0.7531.pt',
                            model=model, optimizer=optimizer,
                            scheduler=scheduler)

    print('Checkpoint learning rate:', tu.get_lr(optimizer))
    print('Start from epoch:', start_epoch)

Output ::

    Original learning rate: 0.001
    Checkpoint learning rate: 0.1234
    Start epoch: 1

.. _statistics:

Statistics:
^^^^^^^^^^^

.. code:: python

    import torch
    import torchutils as tu

    # define your dataset and dataloader
    dataset = MyDataset()
    trainloader = torch.utils.data.DataLoader(dataset, batch_size=1,
                                              num_workers=1,
                                              shuffle=False)

    # get statistics
    stats = tu.get_dataset_stats(trainloader, verbose=True)
    print('Mean:', stats['mean'])
    print('Std:', stats['std'])

Output ::

    Calculating dataset stats...
    Batch 100/100
    Mean: tensor([10000.0098,  9999.9795,  9999.9893])
    Std: tensor([0.9969, 1.0003, 0.9972])

.. _`learning rate`:

Learning Rate:
^^^^^^^^^^^^^^

.. code:: python

    import torchvision
    import torchutils as tu
    import torch.optim as optim

    model = torchvision.models.alexnet()
    optimizer = optim.Adam(model.parameters())

    # get learning rate
    current_lr = tu.get_lr(optimizer)
    print('Current learning rate:', current_lr)

    # set learning rate
    optimizer = tu.set_lr(optimizer, current_lr*0.1)
    revised_lr = tu.get_lr(optimizer)
    print('Revised learning rate:', revised_lr)

Output ::

    Current learning rate: 0.001
    Revised learning rate: 0.0001

.. _`evaluation metrics`:

Evaluation Metrics:
^^^^^^^^^^^^^^^^^^^

.. code:: python

    import torch
    import torch.nn as nn
    import torch.optim as optim
    import torchvision
    import torchvision.transforms as transforms
    import torchutils as tu

    # define your network
    model = MyNet()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters())
    trainset = torchvision.datasets.MNIST(root='./data/', train=True,
                                        download=True,
                                        transform=transforms.ToTensor())
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=60,
                                            shuffle=True, num_workers=2,
                                            drop_last=True)
    n_epochs = 1
    model.train()
    for epoch in range(n_epochs):
        print('Epoch: %d/%d' % (epoch + 1, n_epochs))
        # define loss tracker
        loss_tracker = tu.RunningLoss()
        for batch_idx, (data, target) in enumerate(trainloader):
            optimizer.zero_grad()
            outputs = model(data)
            loss = criterion(outputs, target)
            # update loss tracker with latest loss
            loss_tracker.update(loss.item())
            loss.backward()
            optimizer.step()
            if batch_idx % 100 == 0:
                # easily print latest and average loss
                print(loss_tracker)

Output ::

    Epoch: 1/1
    Loss - Val: 2.2921 Avg: 2.2921
    Loss - Val: 0.5084 Avg: 0.9639
    Loss - Val: 0.6027 Avg: 0.6588
    Loss - Val: 0.1817 Avg: 0.5255
    Loss - Val: 0.1005 Avg: 0.4493
    Loss - Val: 0.2982 Avg: 0.3984
    Loss - Val: 0.3103 Avg: 0.3615
    Loss - Val: 0.0940 Avg: 0.3296
    Loss - Val: 0.0957 Avg: 0.3071
    Loss - Val: 0.0229 Avg: 0.2875

.. _`model summary`:

Model Summary:
^^^^^^^^^^^^^^

.. code:: python

    import torch
    import torchvision
    import torchutils as tu

    model = torchvision.models.alexnet()
    # easily print model summary
    tu.get_model_summary(model, torch.rand((1, 3, 224, 224)))

Output ::

    =========================================================================================
    Layer                           Kernel             Output          Params           FLOPs
    =========================================================================================
    0_features.Conv2d_0         [3, 64, 11, 11]    [1, 64, 55, 55]       23,296    70,470,400
    1_features.ReLU_1                         -    [1, 64, 55, 55]            0             0
    2_features.MaxPool2d_2                    -    [1, 64, 27, 27]            0             0
    3_features.Conv2d_3         [64, 192, 5, 5]   [1, 192, 27, 27]      307,392   224,088,768
    4_features.ReLU_4                         -   [1, 192, 27, 27]            0             0
    5_features.MaxPool2d_5                    -   [1, 192, 13, 13]            0             0
    6_features.Conv2d_6        [192, 384, 3, 3]   [1, 384, 13, 13]      663,936   112,205,184
    7_features.ReLU_7                         -   [1, 384, 13, 13]            0             0
    8_features.Conv2d_8        [384, 256, 3, 3]   [1, 256, 13, 13]      884,992   149,563,648
    9_features.ReLU_9                         -   [1, 256, 13, 13]            0             0
    10_features.Conv2d_10      [256, 256, 3, 3]   [1, 256, 13, 13]      590,080    99,723,520
    11_features.ReLU_11                       -   [1, 256, 13, 13]            0             0
    12_features.MaxPool2d_12                  -     [1, 256, 6, 6]            0             0
    13_classifier.Dropout_0                   -          [1, 9216]            0             0
    14_classifier.Linear_1         [9216, 4096]          [1, 4096]   37,752,832    75,493,376
    15_classifier.ReLU_2                      -          [1, 4096]            0             0
    16_classifier.Dropout_3                   -          [1, 4096]            0             0
    17_classifier.Linear_4         [4096, 4096]          [1, 4096]   16,781,312    33,550,336
    18_classifier.ReLU_5                      -          [1, 4096]            0             0
    19_classifier.Linear_6         [4096, 1000]          [1, 1000]    4,097,000     8,191,000
    =========================================================================================
    Total params: 61,100,840
    Trainable params: 61,100,840
    Non-trainable params: 0
    Total FLOPs: 773,286,232 / 773.29 MFLOPs
    -----------------------------------------------------------------------------------------
    Input size (MB): 0.57
    Forward/backward pass size (MB): 8.31
    Params size (MB): 233.08
    Estimated Total Size (MB): 241.96
    =========================================================================================

.. _`model FLOPs`:

Model FLOPs:
^^^^^^^^^^^^

.. code:: python

    import torch
    import torchvision
    import torchutils as tu

    model = torchvision.models.alexnet()
    # calculate model FLOPs
    total_flops = tu.get_model_flops(model, torch.rand((1, 3, 224, 224)))
    print('Total model FLOPs: {:,}'.format(total_flops))

Output ::

    Total model FLOPs: 773,304,664

.. _`model parameters`:

Model Parameters:
^^^^^^^^^^^^^^^^^

.. code:: python

    import torchvision
    import torchutils as tu

    model = torchvision.models.alexnet()
    # calculate total model parameters
    total_params = tu.get_model_param_count(model)
    print('Total model params: {:,}'.format(total_params))

Output ::

    Total model params: 61,100,840

.. _`random seed`:

Random Seed:
^^^^^^^^^^^^

.. code:: python

    import torchutils as tu

    # set numpy, torch and cuda seed
    tu.set_random_seed(2222)

.. _`gradient flow`:

Gradient Flow:
^^^^^^^^^^^^^^

.. code:: python

    import torch
    import torchvision
    import torchutils as tu

    criterion = torch.nn.CrossEntropyLoss()
    net = torchvision.models.alexnet(num_classes=10)
    out = net(torch.rand(1, 3, 224, 224))
    ground_truth = torch.randint(0, 10, (1, ))
    loss = criterion(out, ground_truth)
    loss.backward()

    # save model gradient flow to image
    tu.plot_gradients(net, './grad_figures/grad_01.png', plot_type='line')

Saved File

.. image:: https://raw.githubusercontent.com/anjandeepsahni/torchutils/master/docs/_static/example_gradient_flow.png
    :alt: Example Gradient Flow 

License
-------
TorchUtils is distributed under the MIT license, see LICENSE.