PyTorch utility APIs.
Project description
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
Latest tests use the following package versions:
Numpy == 1.17.4
PyTorch == 1.4.0
Matplotlib == 3.1.1
Installation
$ pip install torchutils
Documentation
Detailed API documentation is available here.
Examples
Checkpoint:
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 model_20190814-212442_e0_0.7531.pt 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:
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) 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:
import torchvision import torchutils as tu import torch.optim as optim model = torchvision.models.alexnet() optimizer = optim.Adam(model.parameters()) current_lr = tu.get_lr(optimizer) print('Current learning rate:', current_lr) 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:
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)) loss_tracker = tu.RunningLoss() for batch_idx, (data, target) in enumerate(trainloader): optimizer.zero_grad() outputs = model(data) loss = criterion(outputs, target) loss_tracker.update(loss.item()) loss.backward() optimizer.step() if batch_idx % 100 == 0: 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:
import torch import torchvision import torchutils as tu model = torchvision.models.alexnet() 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:
import torch import torchvision import torchutils as tu model = torchvision.models.alexnet() 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:
import torchvision import torchutils as tu model = torchvision.models.alexnet() total_params = tu.get_model_param_count(model) print('Total model params: {:,}'.format(total_params)) # Output Total model params: 61,100,840
Random Seed:
import torchutils as tu tu.set_random_seed(2222)
Gradient Flow:
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() tu.plot_gradients(net, './grad_figures/grad_01.png', plot_type='line') # Saved file
License
TorchUtils is distributed under the MIT license, see LICENSE.
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for torchutils-0.0.3-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | f222898b7c6e940dfca939468b752699ed20e9d704ba71458be2fe8f5776d3d0 |
|
MD5 | edad126d5cd3d5a551309296d27c5c8d |
|
BLAKE2b-256 | 1e644976e99710312e1e39992fc539e2e9e034c9185f080c56a0758f88148357 |