Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
# Installs 375.66 for CUDA 8 on Ubuntu 16.04
wget http://us.download.nvidia.com/tesla/375.66/nvidia-diag-driver-local-repo-ubuntu1604_375.66-1_amd64.deb
sudo dpkg -i nvidia-diag-driver-local-repo-ubuntu1604_375.66-1_amd64.deb
sudo apt-get update
sudo apt-get --allow-unauthenticated --assume-yes install cuda-drivers
sudo reboot now
# Installs 384.66 for CUDA 8 on Ubuntu 16.04
wget http://us.download.nvidia.com/tesla/384.66/nvidia-diag-driver-local-repo-ubuntu1604-384.66_1.0-1_amd64.deb
sudo dpkg -i nvidia-diag-driver-local-repo-ubuntu1604-384.66_1.0-1_amd64.deb
sudo apt-get update
sudo apt-get --allow-unauthenticated --assume-yes install cuda-drivers
sudo reboot now
import resource
import torch
import torch.nn as nn
from torch.autograd import Variable
torch.backends.cudnn.benchmark = True
layers = [
nn.Conv2d(3, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 3, kernel_size=3, padding=1)
]
model = nn.Sequential(*layers).cuda()
loss_fn = nn.L1Loss().cuda()
num_iterations = 10000
N, C, H, W = 64, 3, 64, 64
for t in range(num_iterations):
memory_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
print('Iteration %d, using %.2f MB of memory'
% (t, memory_kb / 1024))
x = Variable(torch.randn(N, C, H, W).cuda())
y = Variable(torch.randn(N, C, H, W).cuda())
y_pred = model(x)
loss = loss_fn(y_pred, y)
loss.backward()
model.zero_grad()
# Install PyTorch 0.2 with pip
sudo apt-get --assume-yes install python3-venv
python3 -m venv env
source env/bin/activate
pip install http://download.pytorch.org/whl/cu80/torch-0.2.0.post1-cp35-cp35m-manylinux1_x86_64.whl
# Run the test script
python leak.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment