Skip to content

Instantly share code, notes, and snippets.

yaroslavvb /
Last active Jan 26, 2020
An example of turning contraction order into sequence of einsum calls
# An example of turning contraction order into sequence of einsum calls
from opt_einsum import helpers as oe_helpers
import opt_einsum as oe
def print_schedule(path, indices, output_subscript, terms):
path: contraction path in einsum optimizer format, ie, [(0,), (2,), (1, 3), (0, 2), (0, 1)]
def index_reduce(values, indices, dim):
"""Reduce values by selecting a single element in dimension dim:
Example below produces rank-2 tensor out of rank-3 values tensor by indexing as follows
dim=0: values[index[i,j],i,j]
dim=1: values[i,index[i,j],j]
dim=2: values[i,j,index[i,j]]
When all entries of "indices" are equal to p, the result is equivalent to slicing along that dimension.
dim=0: values[p,:,:]
View gist:94c78e1d8a6fc45e70fc94b2ba1d4078
import os, numpy as np, random, torch
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.cuda.manual_seed_all(666+int(os.environ.get("RANK", "0")))
def install_pdb_handler():
"""Signals to automatically start pdb:
1. CTRL+\\ breaks into pdb.
2. pdb gets launched on exception.
import signal
import pdb
def handler(_signum, _frame):
yaroslavvb / gist:e53c83c40c8385cd90cdc15c7c61fa63
Last active Sep 23, 2019
Example of Python multi-threading giving a mix of .grad from different backward calls
View gist:e53c83c40c8385cd90cdc15c7c61fa63
import time
import threading
import torch
import torch.nn as nn
def simple_model(d, n):
"""Creates linear neural network initialized to identity"""
layers = []
d = 3
q = torch.rand(d)
q /= q.sum() # empirical distribution (soft labels)
theta = torch.randn(d, requires_grad=True)
p = F.softmax(theta, dim=0)
loss = -torch.sum(q*torch.log(p))
g = p - q
H = torch.diag(p) - outer(p)
torch.allclose(H, hessian(loss, theta))
yaroslavvb /
Created Aug 25, 2019
Example of computing Hessian of linear layer
def test():
data_width = 3
targets_width = 2
batch_size = 3
dataset = TinyMNIST('/tmp', download=True, data_width=data_width, targets_width=targets_width, dataset_size=batch_size)
trainloader =, batch_size=batch_size, shuffle=False)
d1 = data_width ** 2 # hidden layer size, visible size, output size
yaroslavvb /
Created Dec 3, 2017
PyTorch tied autoencoder with l-BFGS
import util as u
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
import numpy as np
# todo: make images global
yaroslavvb /
Created Oct 31, 2017
Example of getting memory allocation stats out of timeline
import tensorflow as tf
import memory_util
import numpy as np
def memory_timeline_from_nodestats(nodestats):
lines = []
for node in nodestats:
mem = node.memory
assert(len(mem) == 1), str(node)
import torch
import numpy as np
import time
def benchmark():
num_iters = 100
def compute():
a = torch.ones(int(SIZE_GB*1000*250000)).cuda()