I have trained a small Flux.jl neural network that I am embedding into a larger model but I need to use the neural network to make millions (billions?) of predictions as part of running this larger model. Since the larger model runs on the GPU I also want to evaluate the neural network on the GPU. …

You want to input a batch into the neural network, so your input is y=CUDA.ones(32, N), where N is the number of inputs you want to process in parallel. This will be the easiest way to parallelise the execution. You should get an output matrix which is 31 by N.

How to efficiently evaluate a Flux.jl neural network millions of times on the GPU?

Specific Domains Machine Learning

PolarizedPoutine October 17, 2022, 7:32pm 3

Thanks for pointing this out @jmair! Can’t believe I didn’t know that you can batch evaluate this easily.

2 Likes

Topic		Replies	Views
Neural network in Flux.jl using CUDA is slower General Usage	0	488	July 15, 2020
Allocation of Memory while evaluate a model Machine Learning flux	7	640	November 8, 2021
DiffEqFlux neural_ode used with Flux.Train! is slower on GPU than CPU GPU flux , ode	0	737	December 31, 2019
Flux: Hard to use train! and DataLoader for minibatched NamedTuple dataset with GPU Machine Learning flux	2	1448	September 24, 2020
Parallelizing Loss Function Machine Learning parallel , multithreading , flux , distributed	4	534	March 15, 2021

How to efficiently evaluate a Flux.jl neural network millions of times on the GPU?

Related topics