model.go

package nn

import (
	"go4ml.xyz/base/fu"
	"go4ml.xyz/base/model"
	"go4ml.xyz/base/tables"
	"go4ml.xyz/iokit"
	"go4ml.xyz/nn/mx"
	"go4ml.xyz/zorros"
	"golang.org/x/xerrors"
	"gopkg.in/yaml.v3"
	"io"
)

// default batch size for general nn training
const DefaultBatchSize = 32

/*
Model is a ANN/Gym definition to train network
*/
type Model struct {
	Network   Block
	Optimizer OptimizerConf
	Loss      mx.Loss
	Input     mx.Dimension
	Seed      int
	BatchSize int
	Predicted string
	Context   mx.Context // CPU by default
}

func (e Model) Feed(ds model.Dataset) model.FatModel {
	return func(workout model.Workout) (*model.Report, error) {
		return Train(e, ds, workout, DefaultModelMap)
	}
}

/*
PredictionModel is the FeaturesMapper factory
*/
type PredictionModel struct {
	features       []string
	predicts       string
	symbol, params iokit.Input
	context        mx.Context
}

/*
Features model uses when maps features
the same as Features in the training dataset
*/
func (pm PredictionModel) Features() []string { return pm.features }

/*
Column name model adds to result table when maps features.
By default it's 'Predicted'
*/
func (pm PredictionModel) Predicted() string { return pm.predicts }

/*
Returns new table with all original columns except features
adding one new column with prediction
*/
func (pm PredictionModel) FeaturesMapper(batchSize int) (fm tables.FeaturesMapper, err error) {
	network, err := Load(pm.context, pm.symbol, pm.params, batchSize)
	if err != nil {
		return
	}
	fm = &FeaturesMapper{model: pm, network: network}
	return
}

/*
Gpu changes context of prediction network to gpu enabled
*/
func (pm PredictionModel) Gpu(no ...int) model.PredictionModel {
	pm.context = mx.GPU0
	if len(no) > 0 {
		pm.context = mx.Gpu(no[0])
	}
	return pm
}

/*
FeaturesMapper maps features to prediction
*/
type FeaturesMapper struct {
	model   PredictionModel
	network *Network
}

/*
MapFeature returns new table with all original columns except features
adding one new column with prediction/calculation
*/
func (fm *FeaturesMapper) MapFeatures(t *tables.Table) (r *tables.Table, err error) {
	var input tables.Matrix
	if input, err = t.Matrix(fm.model.features, fm.network.BatchSize); err != nil {
		return
	}
	out := make([]float32, fm.network.Output.Dim().Total())
	outWidth := fm.network.Output.Dim().Total() / fm.network.BatchSize
	if input.Width != fm.network.Input.Dim().Total()/fm.network.BatchSize {
		return nil, xerrors.Errorf("features does not fit network input")
	}
	if t.Len() > fm.network.BatchSize {
		return nil, xerrors.Errorf("batch size does not fit network input")
	}
	fm.network.Forward(input.Features, out)
	return t.Except(fm.model.features...).With(tables.MatrixColumn(out[0:outWidth*t.Len()], t.Len()), fm.model.predicts), nil
}

/*
Close releases all bounded resources
*/
func (fm *FeaturesMapper) Close() error {
	fm.network.Release()
	return nil
}

func ObjectifyModel(c map[string]iokit.Input) (pm model.PredictionModel, err error) {
	var rd io.ReadCloser
	if _, ok := c[ModelPartInfo]; !ok {
		return nil, zorros.New("it's not neural network model")
	}
	if rd, err = c[ModelPartInfo].Open(); err != nil {
		return
	}
	defer rd.Close()
	cf := map[string]interface{}{}
	if err = yaml.NewDecoder(rd).Decode(&cf); err != nil {
		return
	}
	m := PredictionModel{
		symbol:   c[ModelPartSymbol],
		params:   c[ModelPartParams],
		features: fu.Strings(cf["features"]),
		predicts: cf["predicts"].(string),
	}
	return m, nil
}

func Objectify(source iokit.Input, collection ...string) (fm model.GpuPredictionModel, err error) {
	x := fu.Fnzs(fu.Fnzs(collection...), "model")
	m, err := model.Objectify(source, model.ObjectifyMap{x: ObjectifyModel})
	if err != nil {
		return
	}
	return m[x].(model.GpuPredictionModel), nil
}

func LuckyObjectify(source iokit.Input, collection ...string) model.GpuPredictionModel {
	fm, err := Objectify(source, collection...)
	if err != nil {
		panic(err)
	}
	return fm
}