text_plain_output_1.png
from torch import tensor, nn

import torch.nn.functional as F

x_train, y_train, x_valid, y_valid = get_data()
n, m = x_train.shape
c = y_train.max() + 1
nh = 50
(n, m, c, nh)

class Model(nn.Module):
    def __init__(self, ni, nh, no):
        super().__init__()
        self.layers = [nn.Linear(ni, nh), nn.ReLU(), nn.Linear(nh, no)]
    def __call__(self, x):
        for l in self.layers:
            x = l(x)
        return x
model = Model(m, nh, 10)
pred = model(x_train)
pred.shape

def log_softmax(x):
    return (x.exp() / x.exp().sum(-1, keepdim=True)).log()
sm_pred = log_softmax(pred)
(sm_pred.shape, sm_pred[:3])

def nll(inp, targ):
    return -inp[range(targ.shape[0]), targ].mean()
loss = nll(sm_pred, y_train)
loss

def log_softmax(x):
    return x - x.exp().sum(-1, keepdim=True).log()
test_near(nll(log_softmax(pred), y_train), loss)
def log_softmax(x):
    return x - x.logsumexp(-1, keepdim=True)
test_near(nll(log_softmax(pred), y_train), loss)
test_near(F.nll_loss(F.log_softmax(pred, -1), y_train), loss)
test_near(F.cross_entropy(pred, y_train), loss)

loss_func = F.cross_entropy
def acc(out, yb):
    return (torch.argmax(out, -1) == yb).float().mean()

bs = 64
xb, yb = (x_train[:bs], y_train[:bs])
preds = model(xb)
(loss_func(preds, yb), acc(preds, yb))