<jupyter_start><jupyter_text>Load the data<jupyter_code>MNIST_URL = "http://deeplearning.net/data/mnist/mnist.pkl"
data_path = download_data(MNIST_URL, ext='.gz')
data_path
with gzip.open(data_path, "rb") as f: 
    ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding='latin-1')
    (x_train, y_train, x_valid, y_valid) = map(torch.tensor, (x_train, y_train, x_valid, y_valid))
x_train.shape
y_train.shape
x_valid.shape
y_valid.shape
import math
math.sqrt(784)<jupyter_output><empty_output><jupyter_text>So we have the flattened pixels from 60k 28x28 images, giving us 60k x 784 elements. We break that into a training and a validation set of 50k and 10k examples, respectively.   Basic Matrix Multiplication Here's our first basic stab at matrix multiplication. It's going to be very slow, because we're doing the whole thing in python. It's nice for understanding exactly what we're working with, but we'll see very soon that there are some tricks we can use to speed this up a lot.<jupyter_code># export
def basic_matmul(a, b):
    assert a.shape[1] == b.shape[0]
    output = torch.zeros(a.shape[0], b.shape[1]).float()
    for i in range(a.shape[0]):
        for j in range(b.shape[1]):
            for k in range(a.shape[1]):
                output[i,j] += a[i,k] * b[k,j]
    return output                
mat1 = torch.tensor([[1,2],
                     [3,4]]).float()

mat2 = torch.tensor([[5,6], 
                     [7,8], 
                     [9,10]]).float()
try:
    basic_matmul(mat1, mat2)
    raise ValueError # if it doesn't throw an assertion error, we get here
except AssertionError:
    pass
mat3 = torch.tensor([[5,6], 
                     [7,8]]).float()
expected = torch.tensor([
    [19,22],
    [43,50]]).float()
basic_matmul(mat1, mat3)
# export
def allclose(a, b, tol=1e-3): return torch.allclose(a, b, atol=1e-3)
assert allclose(basic_matmul(mat1, mat3), expected)
NUM_HIDDEN = 10
weights = torch.randn(x_train.shape[1], NUM_HIDDEN)
biases = torch.randn(1)
%time t1=basic_matmul(x_train[:5], weights)<jupyter_output>CPU times: user 804 ms, sys: 0 ns, total: 804 ms
Wall time: 802 ms<jupyter_text>Ok, so this is terribly slow. It takes us 924ms to do the matrix multiplication for 5 rows of x_train with a hidden layer size of 10. This data is super small, and yet it still takes almost a full second. That's definitely not gonna get it done! Alas, we need to speed things along.  PyTorch Operations First things first, we can use pytorch's built-in operations to get rid of the innermost loop. Often when numerical operations are implemented in base pytorch, they're implemented in aten, which makes them super fast.<jupyter_code># export
def dot_prod_matmul(a, b):
    assert a.shape[1] == b.shape[0]
    output = torch.zeros(a.shape[0], b.shape[1]).float()
    for i in range(a.shape[0]):
        for j in range(b.shape[1]):
            output[i,j] = (a[i,:] * b[:,j]).sum()
    return output                
assert allclose(dot_prod_matmul(mat1, mat3), expected)
%time t1=dot_prod_matmul(x_train[:5], weights)<jupyter_output>CPU times: user 4 ms, sys: 0 ns, total: 4 ms
Wall time: 2.86 ms<jupyter_text>Wow! So we're down from 924ms to 2.53ms. That means the basic version is...<jupyter_code>924/2.53<jupyter_output><empty_output><jupyter_text>...365 times slower than the raw pytorch version. Not bad!  Broadcasting The reason the code above is so slow is because all of the work is being done in actual python. Python itself is super slow. Any highly performant operations that run in python typically delegate to a faster language. In the case of pytorch, the operation is delegated to a low-level, highly performant language called aten. Vectorized operations are delegated to aten, which speeds them up.Here are the rules of broadcasting ([source](https://jakevdp.github.io/PythonDataScienceHandbook/02.05-computation-on-arrays-broadcasting.html)):* Rule 1: If the two arrays differ in their number of dimensions, the shape of the one with fewer dimensions is padded with ones on its leading (left) side.* Rule 2: If the shape of the two arrays does not match in any dimension, the array with shape equal to 1 in that dimension is stretched to match the other shape.* Rule 3: If in any dimension the sizes disagree and neither is equal to 1, an error is raised.Here are some examples.<jupyter_code>a = 1; b = torch.tensor([1,2,3])
a + b<jupyter_output><empty_output><jupyter_text>In this case, the scalar is converted to a 0-dim vector, and then an axis of length 1 is prepended onto it making it a 1-dim vector of length 1. Then, it is stretched along that vector to match the size of of the second vector, and the two are added together.<jupyter_code>a = torch.tensor(
    [[1],
     [2],
     [3]]
)
b = torch.tensor([[1,2,3]])
a + b<jupyter_output><empty_output><jupyter_text>In this case, broadcasting happens in both directions. First we look at the vertical dimension. `a` has size 3 while `b` has size 1, so `b` is repeated three times to make them match. Then we look at the horizontal dimension. `b` has size 3 but `a` has size 1, so `a` is repeated three times to match. Then we add them together element wise to get our final answer. We can use broadcasting to remove the second loop of our matmul function. To do this, we'll keep in mind that every row in the output is the product of the corresponding row of the input times the entire second matrix.<jupyter_code># export
def broadcast_mat_mul(a, b): 
    assert a.shape[1] == b.shape[0]
    output = torch.zeros(a.shape[0], b.shape[1]).float()
    for i in range(a.shape[0]):
        # the unsqueeze makes it a column vector, so that the 
        # pointwise multiplication happens across the columns of b. 
        output[i] = (a[i].unsqueeze(-1) * b).sum(dim=0)
    return output                
assert allclose(broadcast_mat_mul(mat1, mat3), expected)
%timeit -n 10 _=broadcast_mat_mul(x_train[:5], weights)<jupyter_output>253 µs ± 29.2 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)<jupyter_text>And so we're down to 253 microseconds -- about 3700 times faster than our original test:<jupyter_code>924 / .253<jupyter_output><empty_output><jupyter_text>Einstein Summation Einstein summation is a technique that was popularized by <NAME> for denoting matrix operations succinctly. When we use einstein summation, we can delegate our operations almost entirely to aten (i.e. removing the first loop entirely) giving us our final speedup.<jupyter_code># export
def matmul(a, b): 
    return torch.einsum('ik,kj->ij', a, b)           
assert allclose(einsum_mat_mul(mat1, mat3), expected)
%timeit -n 10 _=einsum_mat_mul(x_train[:5], weights)<jupyter_output>45.8 µs ± 10.6 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)<jupyter_text>Nice! Now we're down to about 54 micro-seconds, or about 17k times faster. For kicks, let's see if that matches the standard pytorch implementation:<jupyter_code>%timeit -n 10 _=x_train[:5].matmul(weights)<jupyter_output>12.2 µs ± 7.1 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)<jupyter_text>We see that the actual pytorch implementation is still about twice as fast as the best we could do implementing our own, and about 50k times faster than our original.  An aside on BLAS Why is it that the pytorch matmul is about 3x faster than the one we can achieve with raw pytorch? If you think about how matrix multiplication works, go through all of the rows in the first matrix and dot product them with the first column of the second matrix, then we repeat the same thing again for all the rows of the first matrix for the second column, etc. For large matrices, the whole thing does not fit into your CPU cache, and so each time we revisit a row from the first matrix for a new column in the second matrix, we have to retrieve the data from RAM again. This overhead from redundant fetches to memory can be optimized by instead breaking your matrix into smaller matrices and doing the operations one at a time, so that we can fit our data into the CPU cache.The issue for developers is that these low-level linear algebra operations are typically delegated to a Basic Linear Algebra Subprogram or BLAS. These are typically proprietary and written in assembly language by companies like Intel and Nvidia. As such, it's pretty difficult to hack on things at this level as a developer. This is one of the challenges of mathematical programming currently.   Forward Pass Next, we'll be using the matrix multiplication operation we've been studying to build a forward pass framework. First, we'll start with a `linear` function for combining two matrices and a bias.<jupyter_code>def lin(w, b, x): return x@w + b
def relu(x): return torch.clamp_min(x, 0.)
relu(lin(weights, biases, x_train))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08796454142516195, 'fraction_numerical': 0.02375269917036027, 'mean_word_length': 4.122235157159488, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 42, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '624925', 'license': 'Apache-2.0', 'n_tokens_mistral': 2862, 'n_tokens_neox': 2658, 'n_words': 1187, 'path': 'nbs/01_basic_network.ipynb', 'repo_name': 'rforgione/dl_foundations_study', 'size': '8808'}