Merge branch 'binaryclassifier' of https://github.com/tiemvanderdeure…

…/MLJFlux.jl into tiemvanderdeure-binaryclassifier

docs/src/interface/Classification.md

@@ -1,3 +1,7 @@
 ```@docs
 MLJFlux.NeuralNetworkClassifier
+```
+
+```@docs
+MLJFlux.NeuralNetworkBinaryClassifier
 ```

docs/src/interface/Summary.md

@@ -12,6 +12,7 @@ Model Type | Prediction type | `scitype(X) <: _` | `scitype(y) <: _`
 `NeuralNetworkRegressor` | `Deterministic` | `Table(Continuous)` with `n_in` columns | `AbstractVector{<:Continuous)` (`n_out = 1`)
 `MultitargetNeuralNetworkRegressor` | `Deterministic` | `Table(Continuous)` with `n_in` columns | `<: Table(Continuous)` with `n_out` columns
 `NeuralNetworkClassifier` | `Probabilistic` | `<:Table(Continuous)` with `n_in` columns | `AbstractVector{<:Finite}` with `n_out` classes
+`NeuralNetworkBinaryClassifier` | `Probabilistic` | `<:Table(Continuous)` with `n_in` columns | `AbstractVector{<:Finite{2}}` (`n_out = 2`)
 `ImageClassifier` | `Probabilistic` | `AbstractVector(<:Image{W,H})` with `n_in = (W, H)` | `AbstractVector{<:Finite}` with `n_out` classes
 
 

src/MLJFlux.jl

@@ -29,7 +29,7 @@ include("image.jl")
 include("mlj_model_interface.jl")
 
 export NeuralNetworkRegressor, MultitargetNeuralNetworkRegressor
-export NeuralNetworkClassifier, ImageClassifier
+export NeuralNetworkClassifier, NeuralNetworkBinaryClassifier, ImageClassifier
 export CUDALibs, CPU1
 
 include("deprecated.jl")

src/classifier.jl

@@ -14,14 +14,14 @@ function MLJFlux.shape(model::NeuralNetworkClassifier, X, y)
 end
 
 # builds the end-to-end Flux chain needed, given the `model` and `shape`:
-MLJFlux.build(model::NeuralNetworkClassifier, rng, shape) =
+MLJFlux.build(model::Union{NeuralNetworkClassifier, NeuralNetworkBinaryClassifier}, rng, shape) =
     Flux.Chain(build(model.builder, rng, shape...),
         model.finaliser)
 
 # returns the model `fitresult` (see "Adding Models for General Use"
 # section of the MLJ manual) which must always have the form `(chain,
 # metadata)`, where `metadata` is anything extra needed by `predict`:
-MLJFlux.fitresult(model::NeuralNetworkClassifier, chain, y) =
+MLJFlux.fitresult(model::Union{NeuralNetworkClassifier, NeuralNetworkBinaryClassifier}, chain, y) =
     (chain, MLJModelInterface.classes(y[1]))
 
 function MLJModelInterface.predict(model::NeuralNetworkClassifier,
@@ -37,3 +37,25 @@ MLJModelInterface.metadata_model(NeuralNetworkClassifier,
     input=Union{AbstractMatrix{Continuous},Table(Continuous)},
     target=AbstractVector{<:Finite},
     path="MLJFlux.NeuralNetworkClassifier")
+
+#### Binary Classifier
+
+function MLJFlux.shape(model::NeuralNetworkBinaryClassifier, X, y)
+    X = X isa Matrix ? Tables.table(X) : X
+    n_input = Tables.schema(X).names |> length
+    return (n_input, 1) # n_output is always 1 for a binary classifier
+end
+
+function MLJModelInterface.predict(model::NeuralNetworkBinaryClassifier,
+    fitresult,
+    Xnew)
+    chain, levels = fitresult
+    X = reformat(Xnew)
+    probs = vec(chain(X))
+    return MLJModelInterface.UnivariateFinite(levels, probs; augment = true)
+end
+
+MLJModelInterface.metadata_model(NeuralNetworkBinaryClassifier,
+    input=Union{AbstractMatrix{Continuous},Table(Continuous)},
+    target=AbstractVector{<:Finite{2}},
+    path="MLJFlux.NeuralNetworkBinaryClassifier")

src/core.jl

@@ -274,3 +274,9 @@ function collate(model, X, y)
     ymatrix = reformat(y)
     return [_get(Xmatrix, b) for b in row_batches], [_get(ymatrix, b) for b in row_batches]
 end
+function collate(model::NeuralNetworkBinaryClassifier, X, y)
+    row_batches = Base.Iterators.partition(1:nrows(y), model.batch_size)
+    Xmatrix = reformat(X)
+    yvec = (y .== classes(y)[2])' # convert to boolean
+    return [_get(Xmatrix, b) for b in row_batches], [_get(yvec, b) for b in row_batches]
+end