feat: migrate imagenet example to reactant

Author: avik-pal

examples/ImageNet/.gitignore

@@ -0,0 +1 @@
+tiny-imagenet-200*

examples/ImageNet/Project.toml

@@ -1,29 +1,29 @@
 [deps]
 Boltz = "4544d5e4-abc5-4dea-817f-29e4c205d9c8"
-Comonicon = "863f3e99-da2a-4334-8734-de3dacbe5542"
 DataAugmentation = "88a5189c-e7ff-4f85-ac6b-e6158070f02e"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 Format = "1fa38f19-a742-5d3f-a2b9-30dd87b9d5f8"
 ImageIO = "82e4d734-157c-48bb-816b-45c225c6df19"
 ImageMagick = "6218d12a-5da1-5696-b52f-db25d2ecc6d1"
+ImageTransformations = "02fcd773-0e25-5acc-982a-7f6622650795"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
-LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
-MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
-MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
-NCCL = "3fe64909-d7a1-4096-9b7d-7a0f12cf0f6b"
+OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 ParameterSchedulers = "d7d3b36b-41b8-4d0d-a2bf-768c6151755e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+[sources]
+Lux = {path = "../.."}
+
 [compat]
 Boltz = "1"
-Comonicon = "1"
 DataAugmentation = "0.3"
 Dates = "1.10"
 FileIO = "1.16"
@@ -32,17 +32,10 @@ ImageIO = "0.6"
 ImageMagick = "1"
 JLD2 = "0.5.1, 0.6"
 Lux = "1"
-LuxCUDA = "0.3.3"
 MLDataDevices = "1.17"
-MLUtils = "0.4.4"
-MPI = "0.20.21"
-NCCL = "0.1.2"
 OneHotArrays = "0.2.5"
 Optimisers = "0.4.6"
 ParameterSchedulers = "0.4.2"
 Random = "1.10"
 Setfield = "1.1.1"
 Zygote = "0.7"
-
-[sources]
-Lux = {path = "../../"}

examples/ImageNet/main.jl

@@ -34,26 +34,6 @@
 #   --base-path="/home/avik-pal/data/ImageNet/"
 # ```
 
-# ## Package Imports
-
-using Boltz, Lux, MLDataDevices
-## import Metalhead # Install and load this package to use the Metalhead models with Lux
-
-using Dates, Random
-using DataAugmentation,
-    FileIO, MLUtils, OneHotArrays, Optimisers, ParameterSchedulers, Setfield
-using Comonicon, Format
-using JLD2
-using Zygote
-
-using LuxCUDA
-## using AMDGPU # Install and load AMDGPU to train models on AMD GPUs with ROCm
-using MPI: MPI
-## Enables distributed training in Lux. NCCL is needed for CUDA GPUs
-using NCCL: NCCL
-
-const gdev = gpu_device()
-const cdev = cpu_device()
 
 # ## Setup Distributed Training
 
@@ -82,230 +62,9 @@ end
 const is_distributed = total_workers > 1
 const should_log = !is_distributed || local_rank == 0
 
-# ## Data Loading for ImageNet
-
-## We need the data to be in a specific format. See the
-## [README.md](@__REPO_ROOT_URL__/examples/ImageNet/README.md) for more details.
-
-const IMAGENET_CORRUPTED_FILES = [
-    "n01739381_1309.JPEG",
-    "n02077923_14822.JPEG",
-    "n02447366_23489.JPEG",
-    "n02492035_15739.JPEG",
-    "n02747177_10752.JPEG",
-    "n03018349_4028.JPEG",
-    "n03062245_4620.JPEG",
-    "n03347037_9675.JPEG",
-    "n03467068_12171.JPEG",
-    "n03529860_11437.JPEG",
-    "n03544143_17228.JPEG",
-    "n03633091_5218.JPEG",
-    "n03710637_5125.JPEG",
-    "n03961711_5286.JPEG",
-    "n04033995_2932.JPEG",
-    "n04258138_17003.JPEG",
-    "n04264628_27969.JPEG",
-    "n04336792_7448.JPEG",
-    "n04371774_5854.JPEG",
-    "n04596742_4225.JPEG",
-    "n07583066_647.JPEG",
-    "n13037406_4650.JPEG",
-    "n02105855_2933.JPEG",
-    "ILSVRC2012_val_00019877.JPEG",
-]
-
-function load_imagenet1k(base_path::String, split::Symbol)
-    @assert split in (:train, :val)
-    full_path = joinpath(base_path, string(split))
-    synsets = sort(readdir(full_path))
-    @assert length(synsets) == 1000 "There should be 1000 subdirectories in $(full_path)."
-
-    image_files = String[]
-    labels = Int[]
-    for (i, synset) in enumerate(synsets)
-        filenames = readdir(joinpath(full_path, synset))
-        filter!(x -> x ∉ IMAGENET_CORRUPTED_FILES, filenames)
-        paths = joinpath.((full_path,), (synset,), filenames)
-        append!(image_files, paths)
-        append!(labels, repeat([i - 1], length(paths)))
-    end
-
-    return image_files, labels
-end
-
-default_image_size(::Type{Vision.VisionTransformer}, ::Nothing) = 256
-default_image_size(::Type{Vision.VisionTransformer}, size::Int) = size
-default_image_size(_, ::Nothing) = 224
-default_image_size(_, size::Int) = size
-
-struct MakeColoredImage <: DataAugmentation.Transform end
-
-function DataAugmentation.apply(
-    ::MakeColoredImage, item::DataAugmentation.AbstractArrayItem; randstate=nothing
-)
-    data = itemdata(item)
-    (ndims(data) == 2 || size(data, 3) == 1) && (data = cat(data, data, data; dims=Val(3)))
-    return DataAugmentation.setdata(item, data)
-end
-
-struct FileDataset
-    files
-    labels
-    augment
-end
-
-Base.length(dataset::FileDataset) = length(dataset.files)
-
-function Base.getindex(dataset::FileDataset, i::Int)
-    img = Image(FileIO.load(dataset.files[i]))
-    aug_img = itemdata(DataAugmentation.apply(dataset.augment, img))
-    return aug_img, OneHotArrays.onehot(dataset.labels[i], 0:999)
-end
-
-function construct_dataloaders(;
-    base_path::String, train_batchsize, val_batchsize, image_size::Int
-)
-    sensible_println("=> creating dataloaders.")
-
-    train_augment =
-        ScaleFixed((256, 256)) |>
-        Maybe(FlipX(), 0.5) |>
-        RandomResizeCrop((image_size, image_size)) |>
-        PinOrigin() |>
-        ImageToTensor() |>
-        MakeColoredImage() |>
-        ToEltype(Float32) |>
-        Normalize((0.485f0, 0.456f0, 0.406f0), (0.229f0, 0.224f0, 0.225f0))
-    train_files, train_labels = load_imagenet1k(base_path, :train)
-
-    train_dataset = FileDataset(train_files, train_labels, train_augment)
-
-    val_augment =
-        ScaleFixed((image_size, image_size)) |>
-        PinOrigin() |>
-        ImageToTensor() |>
-        MakeColoredImage() |>
-        ToEltype(Float32) |>
-        Normalize((0.485f0, 0.456f0, 0.406f0), (0.229f0, 0.224f0, 0.225f0))
-    val_files, val_labels = load_imagenet1k(base_path, :val)
-
-    val_dataset = FileDataset(val_files, val_labels, val_augment)
-
-    if is_distributed
-        train_dataset = DistributedUtils.DistributedDataContainer(
-            distributed_backend, train_dataset
-        )
-        val_dataset = DistributedUtils.DistributedDataContainer(
-            distributed_backend, val_dataset
-        )
-    end
-
-    train_dataloader = DataLoader(
-        train_dataset;
-        batchsize=train_batchsize ÷ total_workers,
-        partial=false,
-        collate=true,
-        shuffle=true,
-        parallel=true,
-    )
-    val_dataloader = DataLoader(
-        val_dataset;
-        batchsize=val_batchsize ÷ total_workers,
-        partial=true,
-        collate=true,
-        shuffle=false,
-        parallel=true,
-    )
-
-    return gdev(train_dataloader), gdev(val_dataloader)
-end
-
-# ## Model Construction
-
-function construct_model(;
-    rng::AbstractRNG, model_name::String, model_args, pretrained::Bool=false
-)
-    model = getproperty(Vision, Symbol(model_name))(model_args...; pretrained)
-    ps, st = Lux.setup(rng, model) |> gdev
-
-    sensible_println("=> model `$(model_name)` created.")
-    pretrained && sensible_println("==> using pre-trained model`")
-    sensible_println("==> number of trainable parameters: $(Lux.parameterlength(ps))")
-    sensible_println("==> number of states: $(Lux.statelength(st))")
-
-    if is_distributed
-        ps = DistributedUtils.synchronize!!(distributed_backend, ps)
-        st = DistributedUtils.synchronize!!(distributed_backend, st)
-        sensible_println("==> synced model parameters and states across all ranks")
-    end
-
-    return model, ps, st
-end
 
 # ## Optimizer Configuration
 
-function construct_optimizer_and_scheduler(;
-    kind::String,
-    learning_rate::AbstractFloat,
-    nesterov::Bool,
-    momentum::AbstractFloat,
-    weight_decay::AbstractFloat,
-    scheduler_kind::String,
-    cycle_length::Int,
-    damp_factor::AbstractFloat,
-    lr_step_decay::AbstractFloat,
-    lr_step::Vector{Int},
-)
-    sensible_println("=> creating optimizer.")
-
-    kind = Symbol(kind)
-    optimizer = if kind == :adam
-        Adam(learning_rate)
-    elseif kind == :sgd
-        if nesterov
-            Nesterov(learning_rate, momentum)
-        elseif iszero(momentum)
-            Descent(learning_rate)
-        else
-            Momentum(learning_rate, momentum)
-        end
-    else
-        throw(ArgumentError("Unknown value for `optimizer` = $kind. Supported options are: \
-                             `adam` and `sgd`."))
-    end
-
-    optimizer = if iszero(weight_decay)
-        optimizer
-    else
-        OptimiserChain(optimizer, WeightDecay(weight_decay))
-    end
-
-    sensible_println("=> creating scheduler.")
-
-    scheduler_kind = Symbol(scheduler_kind)
-    scheduler = if scheduler_kind == :cosine
-        l0 = learning_rate
-        l1 = learning_rate / 100
-        ComposedSchedule(
-            CosAnneal(l0, l1, cycle_length), Step(l0, damp_factor, cycle_length)
-        )
-    elseif scheduler_kind == :constant
-        Constant(learning_rate)
-    elseif scheduler_kind == :step
-        Step(learning_rate, lr_step_decay, lr_step)
-    else
-        throw(ArgumentError("Unknown value for `lr_scheduler` = $(scheduler_kind). \
-                             Supported options are: `constant`, `step` and `cosine`."))
-    end
-
-    optimizer = if is_distributed
-        DistributedUtils.DistributedOptimizer(distributed_backend, optimizer)
-    else
-        optimizer
-    end
-
-    return optimizer, scheduler
-end
 
 # ## Utility Functions