Add tabular model

src/models/Models.jl

@@ -1,5 +1,6 @@
 module Models
 
+using Base: Bool, Symbol
 using ..FastAI
 
 using BSON
@@ -13,9 +14,11 @@ include("blocks.jl")
 
 include("xresnet.jl")
 include("unet.jl")
+include("tabularmodel.jl")
 
 
-export xresnet18, xresnet50, UNetDynamic
+export xresnet18, xresnet50, UNetDynamic, 
+TabularModel, get_emb_sz, embeddingbackbone, continuousbackbone, classifierbackbone, sigmoidrange
 
 
 end

src/models/tabularmodel.jl

@@ -0,0 +1,81 @@
+function emb_sz_rule(n_cat)
+    min(600, round(1.6 * n_cat^0.56))
+end
+
+function _one_emb_sz(catdict, catcol::Symbol, sz_dict=nothing)
+    sz_dict = isnothing(sz_dict) ? Dict() : sz_dict
+    n_cat = length(catdict[catcol])
+    sz = catcol in keys(sz_dict) ? sz_dict[catcol] : emb_sz_rule(n_cat)
+    Int64(n_cat)+1, Int64(sz)
+end
+
+function get_emb_sz(catdict, cols; sz_dict=nothing)
+    [_one_emb_sz(catdict, catcol, sz_dict) for catcol in cols]
+end
+
+function sigmoidrange(x, low, high)
+    @. Flux.sigmoid(x) * (high - low) + low
+end
+
+function embeddingbackbone(embedding_sizes, dropoutprob=0.)
+    embedslist = [Flux.Embedding(ni, nf) for (ni, nf) in embedding_sizes]
+    emb_drop = dropoutprob==0. ? identity : Dropout(dropoutprob)
+    Chain(
+        x -> tuple(eachrow(x)...), 
+        Parallel(vcat, embedslist), 
+        emb_drop
+    )
+end
+
+function continuousbackbone(n_cont)
+    n_cont > 0 ? BatchNorm(n_cont) : identity
+end
+
+function classifierbackbone(
+        layers;
+        ps=0,
+        use_bn=true,
+        bn_final=false,
+        act_cls=Flux.relu,
+        lin_first=true)
+    ps = Iterators.cycle(ps)
+    classifiers = []
+
+    for (isize, osize, p) in zip(layers[1:(end-1)], layers[2:end], ps)
+        layer = linbndrop(isize, osize; use_bn=use_bn, p=p, act=act_cls, lin_first=lin_first)
+        push!(classifiers, layer)
+    end
+    Chain(classifiers...)
+end
+
+function TabularModel(
+        catbackbone,
+        contbackbone,
+        classifierbackbone; 
+        final_activation=identity)
+    tabularbackbone = Parallel(vcat, catbackbone, contbackbone)
+    Chain(
+        tabularbackbone,
+        classifierbackbone,
+        final_activation
+    )
+end
+
+function TabularModel(
+        catcols,
+        n_cont::Number,
+        out_sz::Number,
+        layers=[200, 100];
+        catdict,
+        sz_dict=nothing,
+        ps=0.)
+    embedszs = get_emb_sz(catdict, catcols, sz_dict=sz_dict)
+    catback = embeddingbackbone(embedszs)
+    contback = continuousbackbone(n_cont)
+
+    classifierin = mapreduce(layer -> size(layer.weight)[1], +, catback[2].layers, init = n_cont)
+    layers = append!([classifierin], layers, [out_sz])
+    classback = classifierbackbone(layers, ps=ps)
+
+    TabularModel(catback, contback, classback)
+end

test/imports.jl

@@ -6,6 +6,7 @@ using FastAI: Image, Keypoints, Mask, testencoding, Label, OneHot, ProjectiveTra
     encodedblock, decodedblock, encode, decode, mockblock
 using FilePathsBase
 using FastAI.Datasets
+using FastAI.Models
 using DLPipelines
 import DataAugmentation
 import DataAugmentation: getbounds

test/models/tabularmodel.jl

@@ -0,0 +1,47 @@
+include("../imports.jl")
+
+@testset ExtendedTestSet "TabularModel Components" begin
+    @testset ExtendedTestSet "embeddingbackbone" begin
+        embed_szs = [(5, 10), (100, 30), (2, 30)]
+        embeds = embeddingbackbone(embed_szs, 0.)
+        x = [rand(1:n) for (n, _) in embed_szs]
+
+        @test size(embeds(x)) == (70, 1)
+    end
+
+    @testset ExtendedTestSet "continuousbackbone" begin
+        n = 5
+        contback = continuousbackbone(n)
+        x = rand(5, 1)
+        @test size(contback(x)) == (5, 1)
+    end
+
+    @testset ExtendedTestSet "classifierbackbone" begin
+        classback = classifierbackbone([10, 200, 100, 2])
+        x = rand(10, 2)
+        @test size(classback(x)) == (2, 2)
+    end
+
+    @testset ExtendedTestSet "TabularModel" begin
+        n = 5
+        embed_szs = [(5, 10), (100, 30), (2, 30)]
+
+        embeds = embeddingbackbone(embed_szs, 0.)
+        contback = continuousbackbone(n)
+        classback = classifierbackbone([75, 200, 100, 4])
+
+        tm = TabularModel(embeds, contback, classback, final_activation = x->FastAI.sigmoidrange(x, 2, 5))
+
+        x = ([rand(1:n) for (n, _) in embed_szs], rand(5, 1))
+        y1 = tm(x)
+        @test size(y1) == (4, 1)
+        @test all(y1.> 2) && all(y1.<5)
+
+        catcols = [:a, :b, :c]
+        catdict = Dict(:a => rand(4), :b => rand(99), :c => rand(1))
+        tm2 = TabularModel(catcols, n, 4, [200, 100], catdict = catdict, sz_dict = Dict(:a=>10, :b=>30, :c=>30))
+        @test size(tm2(x)) == (4, 1)
+    end
+end
+
+

test/runtests.jl

@@ -55,4 +55,10 @@ include("imports.jl")
         end
         # TODO: test learning rate finder
     end
+
+    @testset ExtendedTestSet "models/" begin
+        @testset ExtendedTestSet "tabularmodel.jl" begin
+            include("models/tabularmodel.jl")
+        end
+    end
 end