pytorch_scaling_importance

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()