TorchUtils is a Python package providing helpful utility APIs for your PyTorch projects.
- 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.
- PyTorch >= 1.0.0
- Numpy >= 1.16.2
- Matplotlib >= 3.0.3
PyPi:
$ pip install torchutils
Conda:
$ conda install -c sahni torchutils
Detailed API documentation is available here.
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
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])
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
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
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 =========================================================================================
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
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
import torchutils as tu
# set numpy, torch and cuda seed
tu.set_random_seed(2222)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
TorchUtils is distributed under the MIT license, see LICENSE.