PyTorch

Example Codes

  1. Basic implementation examples (Linear Regression, CNN, ResNet, RNN, GAN, VAE, TensorBoard).

Tensorboard for PyTorch

Save metrics to tensorboard
from tensorboardX import SummaryWriter

tb_logger = SummaryWriter(args.save_path + '/log')
tb_logger.add_scalar('Train/Loss', losses, curr_step)
tb_logger.add_scalar('Train/Top@1', top1.avg, curr_step)
tb_logger.add_scalar('Eval/Loss', val_loss, curr_step)
tb_logger.add_scalar('Eval/Top@1', prec1, curr_step)

tb_logger.flush()
tb_logger.close()

‘model.summary()‘ in PyTorch

VGG16 Output Shape and Param #
import torch
from torchvision import models
from torchsummary import summary

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg = models.vgg16().to(device)

input_shape = (3, 224, 224)
summary(vgg, input_shape)

PyTorch-OpCounter: Counting Flops / MACs

ResNet50 Flops
from torchvision.models import resnet50
from thop import profile

model = resnet50()
input = torch.randn(1, 3, 224, 224)

macs, params = profile(model, inputs=(input, ))

Another Flops counter for CNN

Profile your running time

PyTorch Profiler
import torch
import torchvision.models as models

model = models.densenet121(pretrained=True)
x = torch.randn((1, 3, 224, 224), requires_grad=True)

with torch.autograd.profiler.profile(use_cuda=True) as prof:
    model(x)
print(prof)

Remove randomness in training

Reproducibility for PyTorch
import numpy as np
import random
import os
import torch

def seed_torch(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.enabled = False # this line for *exact* same results.
seed_torch()

Effective PyTorch

PyProf - PyTorch Profiling tool

runx - An experiment management tool