Machine Learning in shaders – 3. Matrix Multiplication

Goal

If we want to build a scalable system for ML in shaders we ought to think about a model where we could define the number of neurons per hidden layer without limitations. With that purpose, we will build a matrix multiplication shader as an independent module that will facilitate the creation of multi-layered Deep Learning models. We will also build a debugger in Python to make sure that our operations are correct.

Matrix multiplication basics

When I write “matrix multiplication” what I really mean is the dot product between two matrices. Let us check out what is the dot product between two vectors so that we can understand the principle:

$begin{bmatrix}a,b,cend{bmatrix}cdotbegin{bmatrix}d\ e\ fend{bmatrix}= (a times d) +(b times e) + (c times f) = scalar$

These are a row and a column vector. We can also see these as (1, 3), and (3,1) matrices. We observe that the dot product of a (1, 3) matrix and a (3, 1) matrix is a (1,1) matrix, which is a single value or scalar. We could predict the resulting sizes of a matrix operation with the simple mnemotechnic we saw in the previous tutorial.

If we take two (2,2) matrices and perform the dot product between them, we have:

Matrix product example from Khan Academy

We can use this principle to predict the size of a (3,2) and (2,1) matrix, which results in a (3, 1) matrix. It is crucial to internalize how matrix multiplication operates since it is single-handedly the most important step to prevent bugs and errors.

Getting matrix sizes right with textures

Say we want to obtain the dot product of a (1, 2) and a (2, 3) matrix. We need to create two textures with those dimensions. Notice that we need to define the textures in what would seem to be the inverse order than how we usually think of matrixes (it is not btw!).

If we want to obtain a (1, 2) looking texture, meaning a matrix of one row and two columns, we need to create a TOP with a resolution of width = 2 and height = 1. That is the same as saying a matrix of (1, 2). “Width” specifies how many columns of pixels there are, and “height” how many rows. It is just that we are used to expressing textures with columns first and rows later, that’s all.

For simplicity, let us build something with random values. Create two shaders and set them to 32-bit float mono, since we only need to use a single value per pixel. Next, set the input and viewer smoothness to “nearest pixel” to disable automatic interpolation and write the following code in each shader:

float random(vec2 st, float seed) {
	return fract(sin(dot(st, vec2(15.63547, 87.84849))) * seed);
}
void main(){    
	vec2 uv = vUV.st;    
	float pixel = random(uv, 9999);    
	vec4 color = vec4(w);    
	fragColor = TDOutputSwizzle(pixel);
}

Now, set one of the shaders (a) to have a resolution of (2×1) and set the other one (b) to (3×2). These represent our (1, 2) and (2, 3) matrices. You can already tell that we expect a (3,1) matrix as output or a texture of resolution (1×3).

Implementation

Let us go ahead and build our matrix multiplication shader. Create a GLSL TOP and connect the operators we just created. This shader needs to be of dimension (width of b, height of a). Create an expression to grab the width and height from each respective shader and plug that into the resolution parameters. Set also the pixel format to 32-bit mono-alpha, since we expect a single value per pixel. I have also disabled the GBA channels using channel masking, so we focus only on the value of a single channel.

Matrix product example from Khan Academy

Now, how to do the required operation?

$$begin{bmatrix}a & bend{bmatrix}cdotbegin{bmatrix}c & d & e\ f & g &  h end{bmatrix}= begin{bmatrix}(ac + bf) (ad+bg) (ae + bh) end{bmatrix}$$

We need to traverse the row of the first matrix and match every element of each column of the second matrix. It sounds like the perfect job for a for loop doesn’t it? So, the most comfortable way to implement this in our shader is to get the coordinates of the texture in pixel values (or texels). We can query both textures across the appropriate axis on each iteration:

vec2 uv = vUV.st; // Get normalized coordinate (0 to 1)
vec2 res = uTDOutputInfo.res.zw; // Get width and height of shader 
ivec2 xy = ivec2(uv*res); // convert coordinates to pixel positions (texel coordinates)

Using the xy variable, we can access the coordinate for the current pixel in the shader. Given that our output texture is (1, 3), notice that when we are in:

$pixel_(0,0): xy.x = 0, xy.y = 0$
$pixel_(1,0): xy.x = 1, xy.y = 0$
$pixel_(2,0): xy.x = 2, xy.y = 0$

If we would access the pixel value of matrix a: $(0, xy.y)$ we are pointing the shader to grab the correct pixel in the column $0$ of $a$, and if we access the pixel value of matrix b: $(xy.x, 0)$, then we are telling the shader to grab the right pixel in the row $0$ of $b$.

Since the rows of $, a$ have matching dimensions with the cols of $b$, we can loop over the size of a row in $a$ to effectively access a whole row and column of values per pixel from each matrix. To do so, we can use the texelFetch() function from GLSL in our loop. We then need to multiply each component of a row and a column and add the results to obtain the final product. For that, we can create a variable to accumulate the sum. Here is the code with all these steps implemented:

uniform int dimX;
out vec4 fragColor;
void main() {
vec2 uv = vUV.st; // Get normalized coordinate (0 to 1)
vec2 res = uTDOutputInfo.res.zw; // Get width and height
ivec2 xy = ivec2(uv*res); // convert coordinates to pixel positions (texels)
float product = 0.0; // declare place holder for final computation
for(int i = 0; i < dimX; i++) {
		float a = texelFetch(sTD2Inputs[0], ivec2(i, xy.y), 0).x; // right col in a
		float b = texelFetch(sTD2Inputs[1], ivec2(xy.x, i), 0).x; // right row in b 
		product += a*b;
	}
	fragColor = vec4(product)
}

There you have it! Your resulting texture should be (3×1), a pixel matrix that contains the result of the operation.

Now, how do we know it is correct? (and do not tell me we have to do it by hand!)

Building a debugger

When uncertainty hits us, the best is to build a tester. We can write a parallel module in Python that uses NumPy’s np. dot() function and compare the results we obtain with our hand-made matrix multiplication shader to be 100% sure that we have done everything right.

I will not explain how to set this up since I am assuming you have proficiency with TouchDesigner and should know how to use select TOPs and DAT Execute operators. It should suffice to say that I am converting the textures we want to multiply to tables, so then I can use a script to gather the pixel values and write the products into another table that then becomes a texture.

Here is the code to do that and a snapshot of the network:

#########################################
#
# Little script to check dot product from
# textures
#
#########################################
import numpy as np
def processTexture(tex):
	out = np.delete(tex, [1,2,3], axis=2)
	return out.reshape(tex.shape[0], tex.shape[1])
# prepare data	
a = processTexture(op(parent().par.A).numpyArray())
b = processTexture(op(parent().par.B).numpyArray())
# take dot product
product = np.dot(a,b)
# Write rows and cols
t = op('table_product')
t.clear()
for i in product:
	t.appendRow(i)
Matrix product example from Khan Academy

Check if the resulting values are correct. Each pixel in both implementations should yield the same results. This would still work with a bigger output texture, say (2, 3), where:

$pixel_(0,0): xy.x = 0, xy.y = 0$
$pixel_(1,0): xy.x = 1, xy.y = 0$
$pixel_(2,0): xy.x = 2, xy.y = 0$
$pixel_(0,1): xy.x = 0, xy.y = 1$
$pixel_(1,1): xy.x = 1, xy.y = 1$
$pixel_(2,1): xy.x = 2, xy.y = 1$

Let’s try and increase the height of a and the width of b to an even bigger number. How about 10? That gives you a (10, 10) matrix. If it does, then we have successfully built our matrix multiplier and can come back to implement more layers in our shallow neural network.

You can find the full project here with an interactive UI to test values randomly or by hand: MatrixMultiplication.tox

Matrix product example from Khan Academy

Comment

This post doesn't have any comment. Be the first one!

hide comments
Follow
...

This is a unique website which will require a more modern browser to work!

Please upgrade today!