Skip to content

Instantly share code, notes, and snippets.

Created December 26, 2018 11:27
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save stared/b598c03ade397baf3fa03c52bd79e90d to your computer and use it in GitHub Desktop.
#!/usr/bin/env python
import delve
import logging
import torch
import torch.nn as nn
from delve import CheckLayerSat
from torch.autograd import Variable
from tqdm import tqdm, trange
class TwoLayerNet(torch.nn.Module):
def __init__(self, D_in, H, D_out):
super(TwoLayerNet, self).__init__()
self.fc = torch.nn.Sequential(
torch.nn.Linear(D_in, H),
torch.nn.Linear(H, D_out)
def forward(self, x):
return self.fc(x)
cuda = torch.cuda.is_available()
for h in [3, 32, 128]:
# N is batch size; D_in is input dimension;
# H is hidden dimension; D_out is output dimension.
N, D_in, H, D_out = 64, 1000, h, 10
# Create random Tensors to hold inputs and outputs
x = Variable(torch.randn(N, D_in))
y = Variable(torch.randn(N, D_out))
model = TwoLayerNet(D_in, H, D_out)
if cuda:
x, y, model = x.cuda(), y.cuda(), model.cuda()
layers = model.parameters()
stats = CheckLayerSat('regression/h{}'.format(h), layers)
loss_fn = torch.nn.MSELoss(size_average=False)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
steps_iter = trange(2000, desc='steps', leave=True, position=0)
"{:^80}".format("Regression - TwoLayerNet - Hidden layer size {}".format(h))
for i in steps_iter:
y_pred = model(x)
loss = loss_fn(y_pred, y)
steps_iter.set_description('loss=%g' %
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment