Skip to content

Instantly share code, notes, and snippets.

@ptrblck
Created May 16, 2018 20:51
Show Gist options
  • Save ptrblck/198886e12b9f6a51fdd81f0f89ae1a7c to your computer and use it in GitHub Desktop.
Save ptrblck/198886e12b9f6a51fdd81f0f89ae1a7c to your computer and use it in GitHub Desktop.
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
class MyDataset(Dataset):
def __init__(self, data, target):
self.data = data
self.target = target
def __getitem__(self, index):
x = self.data[index]
y = self.target[index]
return x, y
def __len__(self):
return len(self.data)
def train():
model.train()
for epoch in range(20):
acc = 0
num_train_samples = 0
for batch_idx, (data, target) in enumerate(loader_train):
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
_, pred = torch.max(output, 1)
acc += (pred==target).float().sum()
num_train_samples += target.size(0)
acc /= num_train_samples
print('Epoch {}, loss {}, acc {}'.format(epoch, loss.item(), acc.item()))
def evaluate():
model.eval()
with torch.no_grad():
acc = 0
num_eval_samples = 0
for batch_idx, (data, target) in enumerate(loader_test):
output = model(data)
_, pred = torch.max(output, 1)
acc += (pred==target).float().sum()
num_eval_samples += target.size(0)
acc /= num_eval_samples
print('Acc {}'.format(acc.item()))
normalize = False
batch_size = 64
n_samples = 1000
useless = torch.empty(n_samples, 5).uniform_(0, 100)
important = (torch.arange(n_samples) / n_samples).unsqueeze(1).repeat(1, 5)
sigma = 0.1
important += torch.randn(n_samples, 5) * 0.1
data = torch.cat((useless, important), 1)
target = torch.cat((torch.zeros(n_samples//2), torch.ones(n_samples//2))).long()
shuffle_idx = torch.randperm(n_samples)
data = data[shuffle_idx]
target = target[shuffle_idx]
data_train = data[:int(0.8*n_samples)]
data_test = data[int(0.8*n_samples):]
target_train = target[:int(0.8*n_samples)]
target_test = target[int(0.8*n_samples):]
# Normalize
if normalize:
from sklearn.preprocessing import StandardScaler
scaler = StandardScaler()
data_train = scaler.fit_transform(data_train.numpy())
data_test = scaler.transform(data_test.numpy())
data_train = torch.from_numpy(data_train)
data_test = torch.from_numpy(data_test)
dataset_train = MyDataset(data_train, target_train)
dataset_test = MyDataset(data_test, target_test)
loader_train = DataLoader(dataset_train,
batch_size=batch_size,
num_workers=0,
shuffle=True)
loader_test = DataLoader(dataset_test,
batch_size=batch_size,
num_workers=0,
shuffle=False)
model = nn.Sequential(
nn.BatchNorm1d(10),
nn.Linear(10, 25),
nn.ReLU(),
nn.Linear(25, 2),
nn.LogSoftmax()
)
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
train()
evaluate()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment