Add PackageExtensionCompat and PrecompileTools

Remove allow_free!
Remove use of Flux.params in src

Project.toml

@@ -12,6 +12,8 @@ MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"
 NaiveNASflux = "85610aed-7d32-5e57-bb50-4c2e1c9e7997"
 NaiveNASlib = "bd45eb3e-47ce-54bd-9eaf-e86c5f900853"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -34,6 +36,8 @@ MemPool = "0.3"
 NaiveNASflux = "2.0.10"
 NaiveNASlib = "2.0.11"
 Optimisers = "0.2"
+PackageExtensionCompat = "1"
+PrecompileTools = "1"
 Reexport = "0.2.0, 1"
 Setfield = "0.3.4, 0.5, 0.6, 0.7, 0.8, 1"
 julia = "1.7"

ext/NaiveGAfluxCUDAExt.jl

@@ -21,6 +21,4 @@ NaiveGAflux.matchdatatype(::NaiveGAfluxCudaDevice, iter) = GpuIterator(iter)
 
 NaiveGAflux._rangetoarr(a::Type{<:CUDA.CuArray}) = a
 
-NaiveGAflux._allow_free!(x::CUDA.CuArray) = CUDA.unsafe_free!(x)
-
 end

src/NaiveGAflux.jl

@@ -10,7 +10,7 @@ using NaiveNASlib.Advanced, NaiveNASlib.Extend
 import Flux
 using Flux: Dense, Conv, ConvTranspose, DepthwiseConv, CrossCor, LayerNorm, BatchNorm, InstanceNorm, GroupNorm, 
             MaxPool, MeanPool, Dropout, AlphaDropout, GlobalMaxPool, GlobalMeanPool, cpu, gpu,
-            SamePad, params
+            SamePad
 import Optimisers
 import Optimisers: WeightDecay
 using Random
@@ -27,7 +27,7 @@ using Setfield
 # For temporary storage of program state for pause/resume type of operations
 using Serialization
 
-const rng_default = MersenneTwister(abs(rand(Int)))
+const rng_default = MersenneTwister(1)
 const modeldir = "models"
 
 # Fitness
@@ -124,4 +124,12 @@ include("iterators.jl")
 include("app/AutoFlux.jl")
 include("visualize/callbacks.jl")
 
+
+using PackageExtensionCompat
+function __init__()
+    @require_extensions
+end
+
+include("precompile.jl")
+
 end # module

src/app/AutoFlux.jl

@@ -24,7 +24,7 @@ Keyword `mdir` is a directory which will be searched for serialized state from w
 function fit(x, y; cb=identity, mdir=missing)
     if ndims(x) == 4 
         outsize = ndims(y) == 1 ? length(unique(y)) : size(y, 1)
-        return fit(ImageClassifier(insize=size(x), outsize, mdir=defaultdir(mdir, "ImageClassifier")), x, y; cb=identity)
+        return fit(ImageClassifier(;insize=size(x), outsize, mdir=modeldir(mdir, "ImageClassifier")), x, y; cb=identity)
     end
     error("No model for $(ndims(x))D data")
 end

src/app/imageclassification/ImageClassification.jl

@@ -1,7 +1,7 @@
 module ImageClassification
 
 using ...NaiveGAflux
-using ..AutoFlux: fit
+using ..AutoFlux: fit, defaultdir
 using NaiveGAflux: GlobalPool
 using NaiveGAflux: shapetrace, squashshapes, fshape, ndimsout, check_apply
 using NaiveGAflux: StatefulGenerationIter

src/batchsize.jl

@@ -303,14 +303,16 @@ function limit_maxbatchsize(bs::ValidationBatchSize,
 end
 
 function maxtrainbatchsize(model, inshape_nobatch, availablebytes=_availablebytes(model))
-    paramsize = mapreduce(ps -> length(ps) * sizeof(eltype(ps)), +, params(model); init=0)
+    elemsize = _model_parsize(model)
+    paramsize = elemsize > 0 ? nparams(model) * elemsize : 0
     actsize = activationsizes(model, inshape_nobatch) 
     den = paramsize + 2 * actsize
     return den > 0 ? fld(availablebytes - paramsize, den) : -1
 end
 
 function maxvalidationbatchsize(model, inshape_nobatch, availablebytes=_availablebytes(model))
-    paramsize = mapreduce(ps -> length(ps) * sizeof(eltype(ps)), +, params(model); init=0)
+    elemsize = _model_parsize(model)
+    paramsize = elemsize > 0 ? nparams(model) * elemsize : 0
     actsize = activationsizes(model, inshape_nobatch)
     return actsize > 0 ? fld(availablebytes - paramsize, actsize) : -1
 end
@@ -330,9 +332,8 @@ function activationsizes(model::CompGraph, inshape_nobatch, elemsize = _model_pa
 end
 
 function _model_parsize(model)
-    ps = params(model)
-    isempty(ps) && return 0
-    return ps |> first |> eltype |> sizeof
+    anyarr = find_first_array(model)
+    anyarr === nothing ? 0 : sizeof(eltype(anyarr))
 end
 
 _availablebytes(model) = _availablebytes(execution_device(model))

src/candidate.jl

@@ -279,7 +279,7 @@ struct FittedCandidate{F, C <: AbstractCandidate} <: AbstractWrappingCandidate
     c::C
 end
 FittedCandidate(c::AbstractCandidate, f::AbstractFitness, gen) = FittedCandidate(gen, fitness(f, c), c)
-FittedCandidate(c::FittedCandidate, f::AbstractFitness, gen) = FittedCandidate(gen, fitness(f, c), wrappedcand(c))
+FittedCandidate(c::FittedCandidate{F}, f::AbstractFitness, gen) where F = FittedCandidate(gen, fitness(f, c)::F, wrappedcand(c))
 
 @functor FittedCandidate
 
@@ -294,8 +294,18 @@ generation(c::FittedCandidate; default=nothing) = c.gen
 # if they are not passed a FittedCandidate. Perhaps having some kind of fitness state container in each candidate?
 newcand(c::FittedCandidate, mapfield) = FittedCandidate(c.gen, c.fitness, newcand(wrappedcand(c), mapfield))
 
+"""
+    nparams(model)
+
+Return the number of trainable parameters in `model`.
+"""
 nparams(c::AbstractCandidate) = model(nparams, c)
-nparams(x) = mapreduce(length, +, params(x).order; init=0)
+nparams(x) = nparams(0, x)
+nparams(x::Integer, g::CompGraph) = nparams(x, vertices(g))
+nparams(x::Integer, v::AbstractVertex) = nparams(x, layer(v))
+nparams(x::Integer, m) = nparams(x, Flux.trainable(m))
+nparams(x::Integer, tr::Union{Tuple, NamedTuple, AbstractArray}) = foldl(nparams, tr; init=x)
+nparams(x::Integer, tr::AbstractArray{<:Number}) = x + length(tr)
 
 """
     MapType{T, F1, F2}

src/fitness.jl

@@ -72,7 +72,6 @@ function _fitness(s::GpuFitness, c::AbstractCandidate)
     fitval = _fitness(s.f, cgpu)
     # In case parameters changed. Would like to do this some other way, perhaps return the candidate too, or move training to evolve...
     transferstate!(c, cpu(cgpu)) # Can't load CuArray into a normal array
-    allow_free!(cgpu) 
     return fitval
 end
 

src/precompile.jl

@@ -0,0 +1,19 @@
+using PrecompileTools
+
+@setup_workload begin
+    @compile_workload begin
+        ## This depends on RNG Implementation which is probably not ideal...
+        Logging.with_logger(Logging.NullLogger()) do
+            AutoFlux.fit(
+                        AutoFlux.ImageClassification.ImageClassifier(;
+                                popsize=2, 
+                                insize=(32,32,3,0),
+                                outsize=10
+                            ), 
+                        zeros(Float32, 32,32,3,0), # Zero size so we at least don't compute and gradients and stuff
+                        zeros(Int, 10, 0); 
+                        stopcriterion = pop -> generation(pop) > 1)
+        end
+    end
+    Random.seed!(rng_default, 1)
+end

src/util.jl

@@ -423,11 +423,3 @@ end
 find_first_array(x::AbstractArray{<:Number}) = x
 
 
-function allow_free!(model) 
-    # fmapstructure to avoid things like trying to create a layer with nothing as weights
-    # We really just want to walk the structure, not map it here
-     Functors.fmapstructure(_allow_free!, model)
-     nothing
-end
-_allow_free!(x) = nothing
-