diff --git a/examples/Basics/main.jl b/examples/Basics/main.jl index dc13e2d92..1dd56eac5 100644 --- a/examples/Basics/main.jl +++ b/examples/Basics/main.jl @@ -116,10 +116,14 @@ W * x # the `cu` function (or the `gpu` function exported by `Lux``), and it supports all of the # above operations with the same syntax. -using LuxCUDA +using LuxCUDA, LuxAMDGPU + if LuxCUDA.functional() x_cu = cu(rand(5, 3)) @show x_cu +elseif LuxAMDGPU.functional() # Similarly, for AMDGPU + x_amd = roc(rand(5, 3)) + @show x_amd end # ## (Im)mutability diff --git a/examples/DDIM/generate.jl b/examples/DDIM/generate.jl index 83002d077..bf8209394 100644 --- a/examples/DDIM/generate.jl +++ b/examples/DDIM/generate.jl @@ -41,7 +41,7 @@ end mkpath(output_dir) - if CUDA.functional() + if LuxCUDA.functional() || LuxAMDGPU.functional() println("GPU is available.") else println("GPU is not available.") diff --git a/examples/GravitationalWaveForm/main.jl b/examples/GravitationalWaveForm/main.jl index 030aae2ae..d8c5169ef 100644 --- a/examples/GravitationalWaveForm/main.jl +++ b/examples/GravitationalWaveForm/main.jl @@ -17,6 +17,7 @@ close(pkg_io) #hide using Lux, ComponentArrays, LineSearches, LuxAMDGPU, LuxCUDA, OrdinaryDiffEq, Optimization, OptimizationOptimJL, Random, SciMLSensitivity using CairoMakie, MakiePublication + CUDA.allowscalar(false) # ## Define some Utility Functions diff --git a/examples/HyperNet/main.jl b/examples/HyperNet/main.jl index 9d34b76f9..22aff09f7 100644 --- a/examples/HyperNet/main.jl +++ b/examples/HyperNet/main.jl @@ -11,6 +11,7 @@ Pkg.precompile(; io=pkg_io) #hide close(pkg_io) #hide using Lux, ComponentArrays, LuxAMDGPU, LuxCUDA, MLDatasets, MLUtils, OneHotArrays, Optimisers, Random, Setfield, Statistics, Zygote + CUDA.allowscalar(false) # ## Loading Datasets diff --git a/examples/ImageNet/Project.toml b/examples/ImageNet/Project.toml index 2494c143b..9221bbf01 100644 --- a/examples/ImageNet/Project.toml +++ b/examples/ImageNet/Project.toml @@ -4,6 +4,7 @@ Boltz = "4544d5e4-abc5-4dea-817f-29e4c205d9c8" Configurations = "5218b696-f38b-4ac9-8b61-a12ec717816d" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" FLoops = "cc61a311-1640-44b5-9fba-1b764f453329" +Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" FluxMPI = "acf642fa-ee0e-513a-b9d5-bcd150f7ec3b" Formatting = "59287772-0a20-5a39-b81b-1366585eb4c0" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" @@ -14,6 +15,7 @@ Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxAMDGPU = "83120cb1-ca15-4f04-bf3b-6967d2e6b60b" LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" +Metalhead = "dbeba491-748d-5e0e-a39e-b530a07fa0cc" OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" diff --git a/examples/ImageNet/main.jl b/examples/ImageNet/main.jl index 19ad46100..25476ebe7 100644 --- a/examples/ImageNet/main.jl +++ b/examples/ImageNet/main.jl @@ -2,6 +2,8 @@ using Augmentor # Image Augmentation using Boltz # Computer Vision Models +import Flux # Some Blotz Models need Flux +import Metalhead # Some Blotz Models need MetalHead using Configurations # Experiment Configurations using LuxAMDGPU # AMDGPUs <3 using LuxCUDA # NVIDIA GPUs <3 @@ -62,7 +64,8 @@ function construct(rng::AbstractRNG, cfg::ModelConfig, ecfg::ExperimentConfig) end function construct(cfg::OptimizerConfig) - if cfg.name == "adam" + opt = Adam(cfg.learning_rate) + if cfg.name == "adam" + opt = Adam(cfg.learning_rate) elseif cfg.name == "sgd" if cfg.nesterov opt = Nesterov(cfg.learning_rate, cfg.momentum) diff --git a/examples/ImageNet/utils.jl b/examples/ImageNet/utils.jl index f3c581ec5..a8363a5a9 100644 --- a/examples/ImageNet/utils.jl +++ b/examples/ImageNet/utils.jl @@ -1,7 +1,13 @@ CUDA.allowscalar(false) -# unsafe_free OneHotArrays -CUDA.unsafe_free!(x::OneHotArray) = CUDA.unsafe_free!(x.indices) + +if LuxCUDA.functional() + # unsafe_free OneHotArrays + CUDA.unsafe_free!(x::OneHotArray) = CUDA.unsafe_free!(x.indices) +elseif LuxAMDGPU.functional() + # unsafe_free OneHotArrays + AMDGPU.unsafe_free!(x::OneHotArray) = AMDGPU.unsafe_free!(x.indices) +end # Loss Function logitcrossentropyloss(ŷ, y) = mean(-sum(y .* logsoftmax(ŷ; dims=1); dims=1)) diff --git a/examples/NeuralODE/main.jl b/examples/NeuralODE/main.jl index 662b9d897..37a419c83 100644 --- a/examples/NeuralODE/main.jl +++ b/examples/NeuralODE/main.jl @@ -18,6 +18,7 @@ using Lux, ComponentArrays, SciMLSensitivity, LuxAMDGPU, LuxCUDA, Optimisers, OrdinaryDiffEq, Random, Statistics, Zygote, OneHotArrays, InteractiveUtils import MLDatasets: MNIST import MLUtils: DataLoader, splitobs + CUDA.allowscalar(false) # ## Loading MNIST diff --git a/examples/SimpleRNN/main.jl b/examples/SimpleRNN/main.jl index cfff2185e..16d8681f6 100644 --- a/examples/SimpleRNN/main.jl +++ b/examples/SimpleRNN/main.jl @@ -84,7 +84,7 @@ function (s::SpiralClassifier)( ## See that the parameters and states are automatically populated into a field called ## `lstm_cell` We use `eachslice` to get the elements in the sequence without copying, ## and `Iterators.peel` to split out the first element for LSTM initialization. - x_init, x_rest = Iterators.peel(eachslice(x; dims=2)) + x_init, x_rest = Iterators.peel(Lux._eachslice(x, Val(2))) (y, carry), st_lstm = s.lstm_cell(x_init, ps.lstm_cell, st.lstm_cell) ## Now that we have the hidden state and memory in `carry` we will pass the input and ## `carry` jointly @@ -119,7 +119,7 @@ function compute_loss(x, y, model, ps, st) return binarycrossentropy(y_pred, y), y_pred, st end -matches(y_pred, y_true) = sum((y_pred .> 0.5) .== y_true) +matches(y_pred, y_true) = sum((y_pred .> 0.5f0) .== y_true) accuracy(y_pred, y_true) = matches(y_pred, y_true) / length(y_pred) # Finally lets create an optimiser given the model parameters.