More work done on torch

logic/models/train/torch/torch.go

@@ -2,14 +2,12 @@ package train
 
 import (
 	types "git.andr3h3nriqu3s.com/andr3/fyp/logic/db_types"
+	my_nn "git.andr3h3nriqu3s.com/andr3/fyp/logic/models/train/torch/nn"
 
+	"git.andr3h3nriqu3s.com/andr3/gotch"
 	"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"
+	torch "git.andr3h3nriqu3s.com/andr3/gotch/ts"
 )
 
 type IForwardable interface {
@@ -18,23 +16,55 @@ type IForwardable interface {
 
 // Container for a model
 type ContainerModel struct {
-	Seq *nn.SequentialT
-	Vs  *nn.VarStore
+	Layers []my_nn.MyLayer
+	Vs     *my_nn.VarStore
+	path   *my_nn.Path
 }
 
 func (n *ContainerModel) ForwardT(x *torch.Tensor, train bool) *torch.Tensor {
-	return n.Seq.ForwardT(x, train)
+	if len(n.Layers) == 0 {
+		return x.MustShallowClone()
+	}
+
+	if len(n.Layers) == 1 {
+		return n.Layers[0].ForwardT(x, train)
+	}
+
+	// forward sequentially
+	outs := make([]*torch.Tensor, len(n.Layers))
+	for i := 0; i < len(n.Layers); i++ {
+		if i == 0 {
+			outs[0] = n.Layers[i].ForwardT(x, train)
+			defer outs[0].MustDrop()
+		} else if i == len(n.Layers)-1 {
+            return n.Layers[i].ForwardT(outs[i-1], train)
+		} else {
+			outs[i] = n.Layers[i].ForwardT(outs[i-1], train)
+			defer outs[i].MustDrop()
+		}
+	}
+    panic("Do not reach here")
 }
 
 func (n *ContainerModel) To(device gotch.Device) {
 	n.Vs.ToDevice(device)
+	for _, layer := range n.Layers {
+		layer.ExtractFromVarstore(n.Vs)
+	}
+}
+
+func (n *ContainerModel) Refresh() {
+	for _, layer := range n.Layers {
+		layer.ExtractFromVarstore(n.Vs)
+	}
 }
 
 func BuildModel(layers []*types.Layer, _lastLinearSize int64, addSigmoid bool) *ContainerModel {
 
-	base_vs := nn.NewVarStore(gotch.CPU)
+	base_vs := my_nn.NewVarStore(gotch.CPU)
 	vs := base_vs.Root()
-	seq := nn.SeqT()
+
+	m_layers := []my_nn.MyLayer{}
 
 	var lastLinearSize int64 = _lastLinearSize
 	lastLinearConv := []int64{}
@@ -46,18 +76,19 @@ func BuildModel(layers []*types.Layer, _lastLinearSize int64, addSigmoid bool) *
 		} 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]))
+			m_layers = append(m_layers, NewLinear(vs, lastLinearSize, shape[0]))
 			lastLinearSize = shape[0]
 		} else if layer.LayerType == types.LAYER_FLATTEN {
-			seq.Add(NewFlatten())
+			m_layers = append(m_layers, NewFlatten())
 			lastLinearSize = 1
 			for _, i := range lastLinearConv {
 				lastLinearSize *= i
 			}
 			log.Info("Flatten: ", "In:", lastLinearConv, "out:", lastLinearSize)
 		} else if layer.LayerType == types.LAYER_SIMPLE_BLOCK {
+			panic("TODO")
 			log.Info("New Block: ", "In:", lastLinearConv, "out:", []int64{lastLinearConv[1] / 2, lastLinearConv[2] / 2, 128})
-			seq.Add(NewSimpleBlock(vs, lastLinearConv[0]))
+			//m_layers = append(m_layers, NewSimpleBlock(vs, lastLinearConv[0]))
 			lastLinearConv[0] = 128
 			lastLinearConv[1] /= 2
 			lastLinearConv[2] /= 2
@@ -65,12 +96,13 @@ func BuildModel(layers []*types.Layer, _lastLinearSize int64, addSigmoid bool) *
 	}
 
 	if addSigmoid {
-		seq.Add(NewSigmoid())
+		m_layers = append(m_layers, NewSigmoid())
 	}
 
 	b := &ContainerModel{
-		Seq: seq,
-		Vs:  base_vs,
+		Layers: m_layers,
+		Vs:     base_vs,
+		path:   vs,
 	}
 	return b
 }