sparse_optimization

sparse_optimization.py

import numpy as np
import scipy.sparse as sp
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset


def one_hot(index, n_cat):
    onehot = torch.zeros(index.size(0), n_cat, device=index.device)
    onehot.scatter_(1, index.type(torch.long), 1)
    return onehot.type(torch.float32)


class MyDataset(Dataset):
    def __init__(self, n_samples, n_input):
        self.X = sp.csr_matrix(np.random.binomial(1, 0.01, size=(n_samples, n_input)))
        self.y = np.random.binomial(1, 0.5, size=(n_samples, 1)).astype(np.int32)
        self.total_size = n_samples

    def __len__(self):
        return self.total_size

    def __getitem__(self, idx):
        return idx

    def collate_fn(self, batch):
        indexes = np.array(batch)
        return torch.FloatTensor(self.X[indexes]), torch.LongTensor(self.y[indexes])


# Neural Network Model (1 hidden layer)
class Net(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out


n_input = 50000
n_samples = 2000
net = Net(n_input, 128, 2)
net.cuda()

my_dataset = MyDataset(n_samples, n_input)
train_loader = DataLoader(my_dataset, batch_size=128, collate_fn=my_dataset.collate_fn, num_workers=1)

# Loss and Optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters())

# Train the Model
for epoch in range(10):
    for i, (x, y) in enumerate(train_loader):
        # Convert torch tensor to Variable
        x = x.cuda()
        y = y.cuda()

        # Forward + Backward + Optimize
        optimizer.zero_grad()  # zero the gradient buffer
        outputs = net(x)
        one_hot(y, 2)
        loss = criterion(outputs, y.view(-1))
        loss.backward()
        optimizer.step()