82 lines
2.0 KiB
Go
82 lines
2.0 KiB
Go
package train
|
|
|
|
import (
|
|
types "git.andr3h3nriqu3s.com/andr3/fyp/logic/db_types"
|
|
|
|
"github.com/charmbracelet/log"
|
|
"github.com/sugarme/gotch"
|
|
"github.com/sugarme/gotch/nn"
|
|
|
|
//"github.com/sugarme/gotch"
|
|
//"github.com/sugarme/gotch/vision"
|
|
torch "github.com/sugarme/gotch/ts"
|
|
)
|
|
|
|
type IForwardable interface {
|
|
Forward(xs *torch.Tensor) *torch.Tensor
|
|
}
|
|
|
|
// Container for a model
|
|
type ContainerModel struct {
|
|
Seq *nn.SequentialT
|
|
Vs *nn.VarStore
|
|
}
|
|
|
|
func (n *ContainerModel) ForwardT(x *torch.Tensor, train bool) *torch.Tensor {
|
|
return n.Seq.ForwardT(x, train)
|
|
}
|
|
|
|
func (n *ContainerModel) To(device gotch.Device) {
|
|
n.Vs.ToDevice(device)
|
|
}
|
|
|
|
func BuildModel(layers []*types.Layer, _lastLinearSize int64, addSigmoid bool) *ContainerModel {
|
|
|
|
base_vs := nn.NewVarStore(gotch.CPU)
|
|
vs := base_vs.Root()
|
|
seq := nn.SeqT()
|
|
|
|
var lastLinearSize int64 = _lastLinearSize
|
|
lastLinearConv := []int64{}
|
|
|
|
for _, layer := range layers {
|
|
if layer.LayerType == types.LAYER_INPUT {
|
|
lastLinearConv = layer.GetShape()
|
|
log.Info("Input: ", "In:", lastLinearConv)
|
|
} else if layer.LayerType == types.LAYER_DENSE {
|
|
shape := layer.GetShape()
|
|
log.Info("New Dense: ", "In:", lastLinearSize, "out:", shape[0])
|
|
seq.Add(NewLinear(vs, lastLinearSize, shape[0]))
|
|
lastLinearSize = shape[0]
|
|
} else if layer.LayerType == types.LAYER_FLATTEN {
|
|
seq.Add(NewFlatten())
|
|
lastLinearSize = 1
|
|
for _, i := range lastLinearConv {
|
|
lastLinearSize *= i
|
|
}
|
|
log.Info("Flatten: ", "In:", lastLinearConv, "out:", lastLinearSize)
|
|
} else if layer.LayerType == types.LAYER_SIMPLE_BLOCK {
|
|
log.Info("New Block: ", "In:", lastLinearConv, "out:", []int64{lastLinearConv[1] / 2, lastLinearConv[2] / 2, 128})
|
|
seq.Add(NewSimpleBlock(vs, lastLinearConv[0]))
|
|
lastLinearConv[0] = 128
|
|
lastLinearConv[1] /= 2
|
|
lastLinearConv[2] /= 2
|
|
}
|
|
}
|
|
|
|
if addSigmoid {
|
|
seq.Add(NewSigmoid())
|
|
}
|
|
|
|
b := &ContainerModel{
|
|
Seq: seq,
|
|
Vs: base_vs,
|
|
}
|
|
return b
|
|
}
|
|
|
|
func SaveModel(model *ContainerModel, modelFn string) (err error) {
|
|
model.Vs.ToDevice(gotch.CPU)
|
|
return model.Vs.Save(modelFn)
|
|
}
|