experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Go: Deallocate large TF_TString on Tensor finalization

Browse Source
This commit is contained in:
Artemiy Ryabinkov 2021-01-19 21:53:46 +02:00 committed by Artemiy Ryabinkov
parent 569095ba3d
commit bd46b2d9ca
1 changed files with 18 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
tensorflow/go/tensor.go
View File

@ -95,8 +95,17 @@ func NewTensor(value interface{}) (*Tensor, error) {
c: C.TF_AllocateTensor(C.TF_DataType(dataType), shapePtr, C.int(len(shape)), C.size_t(nbytes)),
shape: shape,
}
runtime.SetFinalizer(t, (*Tensor).finalize)
raw := tensorData(t.c)
runtime.SetFinalizer(t, func(t *Tensor) {
if dataType == String {
t.clearTStrings(raw, nflattened)
}
t.finalize()
})
buf := bytes.NewBuffer(raw[:0:len(raw)])
if isAllArray(val.Type()) {
@ -206,6 +215,14 @@ func newTensorFromC(c *C.TF_Tensor) *Tensor {
return t
}
func (t *Tensor) clearTStrings(raw []byte, n int64) {
tstrs := (*(*[]C.TF_TString)(unsafe.Pointer(&raw)))[:n]
for _, tstr := range tstrs {
C.TF_TString_Dealloc(&tstr)
}
}
func (t *Tensor) finalize() { C.TF_DeleteTensor(t.c) }
// DataType returns the scalar datatype of the Tensor.