fyp/logic/models/train/torch/torch.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)
}