torchimg

A python package for illustrating pytorch models

How to use

Installation

pip install git+https://github.com/not-lain/torchimg.git

import torch
from torchimg import extract_compute_graph


class SimpleModel(torch.nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = torch.nn.Linear(3 * 224 * 224, 512)
        self.fc2 = torch.nn.Linear(512, 256)
        self.fc3 = torch.nn.Linear(256, 10)

    def forward(self, x):
        x = x.view(-1, 3 * 224 * 224)
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.relu(self.fc2(x))
        x = self.fc3(x)
        return x


model = SimpleModel()
input_tensor = torch.randn(1, 3, 224, 224)
graph = extract_compute_graph(model, input_tensor)