Package: torchutils
TorchUtilsis a Python package providing helpful utility APIs for your
PyTorch projects.FeaturesSave/loadcheckpoints.Calculate datasetstatistics(mean, std, var). Also calculate and track running statistics of data.Get/setlearning rate.Trackevaluation metricssuch as accuracy, running loss, hamming loss.Printmodel summary. Supports: Linear/MLP, Convolution Network, Recurrent Network (RNN/LSTM/GRU), Recursive Network.Calculatemodel FLOPs.Calculate totalmodel parameters.Setrandom seed.Visualizegradient flowin your network.RequirementsPyTorch >= 1.0.0Numpy >= 1.16.2Matplotlib >= 3.0.3InstallationPyPi:$ pip install torchutilsConda:$ conda install -c sahni torchutilsDocumentationDetailed API documentation is availablehere.ExamplesCheckpoint:importtorchvisionimporttorchutilsastuimporttorch.optimasoptimmodel=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 checkpointstart_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)OutputOriginal learning rate: 0.001
Checkpoint learning rate: 0.1234
Start epoch: 1Statistics:importtorchimporttorchutilsastu# define your dataset and dataloaderdataset=MyDataset()trainloader=torch.utils.data.DataLoader(dataset,batch_size=1,num_workers=1,shuffle=False)# get statisticsstats=tu.get_dataset_stats(trainloader,verbose=True)print('Mean:',stats['mean'])print('Std:',stats['std'])OutputCalculating dataset stats...
Batch 100/100
Mean: tensor([10000.0098,  9999.9795,  9999.9893])
Std: tensor([0.9969, 1.0003, 0.9972])Learning Rate:importtorchvisionimporttorchutilsastuimporttorch.optimasoptimmodel=torchvision.models.alexnet()optimizer=optim.Adam(model.parameters())# get learning ratecurrent_lr=tu.get_lr(optimizer)print('Current learning rate:',current_lr)# set learning rateoptimizer=tu.set_lr(optimizer,current_lr*0.1)revised_lr=tu.get_lr(optimizer)print('Revised learning rate:',revised_lr)OutputCurrent learning rate: 0.001
Revised learning rate: 0.0001Evaluation Metrics:importtorchimporttorch.nnasnnimporttorch.optimasoptimimporttorchvisionimporttorchvision.transformsastransformsimporttorchutilsastu# define your networkmodel=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=1model.train()forepochinrange(n_epochs):print('Epoch:%d/%d'%(epoch+1,n_epochs))# define loss trackerloss_tracker=tu.RunningLoss()forbatch_idx,(data,target)inenumerate(trainloader):optimizer.zero_grad()outputs=model(data)loss=criterion(outputs,target)# update loss tracker with latest lossloss_tracker.update(loss.item())loss.backward()optimizer.step()ifbatch_idx%100==0:# easily print latest and average lossprint(loss_tracker)OutputEpoch: 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.2875Model Summary:importtorchimporttorchvisionimporttorchutilsastumodel=torchvision.models.alexnet()# easily print model summarytu.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:importtorchimporttorchvisionimporttorchutilsastumodel=torchvision.models.alexnet()# calculate model FLOPstotal_flops=tu.get_model_flops(model,torch.rand((1,3,224,224)))print('Total model FLOPs:{:,}'.format(total_flops))OutputTotal model FLOPs: 773,304,664Model Parameters:importtorchvisionimporttorchutilsastumodel=torchvision.models.alexnet()# calculate total model parameterstotal_params=tu.get_model_param_count(model)print('Total model params:{:,}'.format(total_params))OutputTotal model params: 61,100,840Random Seed:importtorchutilsastu# set numpy, torch and cuda seedtu.set_random_seed(2222)Gradient Flow:importtorchimporttorchvisionimporttorchutilsastucriterion=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 imagetu.plot_gradients(net,'./grad_figures/grad_01.png',plot_type='line')Saved FileLicenseTorchUtils is distributed under the MIT license, see LICENSE.