Memory leaks

[Joe Bloxsome]

I'm getting memory leaks even after calling MustDrop() on tensors.

This is my code:

package gopt

import (
	"log"
	"net/http"
	"errors"
	"os"

	"github.com/sugarme/gotch"
	"github.com/sugarme/gotch/ts"
	"github.com/sugarme/gotch/vision"
)

func GetFileContentType(path string) (string, error) {
	// Open File
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()

	// Only the first 512 bytes are used to sniff the content type.
	buffer := make([]byte, 512)

	_, err = f.Read(buffer)
	if err != nil {
		return "", err
	}

	// Use the net/http package's handy DectectContentType function. Always returns a valid
	// content-type by returning "application/octet-stream" if no others seemed to match.
	contentType := http.DetectContentType(buffer)

	return contentType, nil
}

type GoPt struct {
	Model  *ts.CModule
	Labels []string
	Imnet  *vision.ImageNet
}

func (gopt *GoPt) LoadModel(path string) {
	// Create ImageNet object to use for input resizing
	gopt.Imnet = vision.NewImageNet()

	model, err := ts.ModuleLoadOnDevice(path, gotch.CPU)
	if err != nil {
		log.Fatal(err)
	}
	gopt.Model = model
}

func (gopt *GoPt) Predict(path string) (string, error) {
	// Check the file is an image first
	contentType, err := GetFileContentType(path)
	if err != nil {
		return "", err
	}

	if contentType != "image/jpeg" {
		return "", errors.New("must be an image file")
	}

	// Load the image file and resize it
	image, err := gopt.Imnet.LoadImageAndResize224(path)
	if err != nil {
		return "", err
	}

	// Apply the forward pass of the model to get the logits.
	unsqueezed := image.MustUnsqueeze(int64(0), false)
	image.MustDrop()
	raw_output := unsqueezed.ApplyCModule(gopt.Model)
	unsqueezed.MustDrop()
	output := raw_output.MustSoftmax(-1, gotch.Float, true)
	raw_output.MustDrop()

	// Convert to list of floats to represent label probabilities
	probs := output.Vals().([]float32)
	output.MustDrop()

	maxVal := probs[0]
	maxIndex := 0
	for i, v := range probs {
		if (v > maxVal) {
			maxVal = v
			maxIndex = i
		}
	}
	maxVal = nil

	return gopt.Labels[maxIndex], nil
}

Whenever I call Predict in a loop, the memory usage is blowing up.

Here's the code calling it:

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"time"

	"github.com/jbloxsome/gopt/gopt"
)

var (
	modelPath string
)

func init() {
	flag.StringVar(&modelPath, "modelpath", "./model.pt", "full path to exported pytorch model.")
}

func main() {
	flag.Parse()

	gopt := gopt.GoPt{
		Labels: []string{
			"false",
			"true",
		},
	}

	gopt.LoadModel(modelPath)

	root := "/test/false"

	files, err := ioutil.ReadDir(root)
	if err != nil {
		log.Fatal(err)
	}

	positive := 0
	negative := 0

	for i, file := range files {
		fmt.Println(i, file.Name())

		pred, err := gopt.Predict(root + "/" + file.Name())
		if err != nil {
			fmt.Println(err)
		}
		if pred == "true" {
			positive = positive + 1
		}
		if pred == "false" {
			negative = negative + 1
		}

		fmt.Println(pred)
	}

	fmt.Println("positives", positive)
	fmt.Println("negatives", negative)
	fmt.Println("total", positive+negative)
	percentPositive := 100 * positive / (positive + negative)
	fmt.Println("positive (%)", percentPositive)
	time.Sleep(time.Second * 60)
}

[sugarme]

@jbloxsome ,

I can't spot any leak in your code, just some double/verbal tensor deletions. Any chance that you can share a testing model so that I can run the code? Thanks.

Also, if you have insight into your model, you can construct model with gotch and load weights from trained model in Pytorch or even train with gotch. Have a look at subpackage https://github.com/sugarme/gotch/tree/master/pickle and example at https://github.com/sugarme/gotch/tree/master/example/pickle .

[Joe Bloxsome]

@sugarme Thanks for taking a look! The model I'm using is https://storage.googleapis.com/freelance-models/model.pt

Ah ok, I'll take a look at that second option as well. The model is resnet32 (fine tuned using FastAI).

[sugarme]

@jbloxsome ,

I tried to run your code with the model and can see memory consuming. It turned out the model was loaded every iteration as you passed Go struct instead of pointer.

Have a look at the code below. It runs just Okay.

package main

import (
	"log"

	"github.com/sugarme/gotch"
	"github.com/sugarme/gotch/ts"
	"github.com/sugarme/gotch/vision"
)

type Model struct {
	cmodule  *ts.CModule
	imageNet *vision.ImageNet
}

func NewModel(modelPath string) (*Model, error) {
	model, err := ts.ModuleLoadOnDevice(modelPath, gotch.CPU)
	if err != nil {
		return nil, err
	}

	imageNet := vision.NewImageNet()

	return &Model{
		cmodule:  model,
		imageNet: imageNet,
	}, nil
}

func (m *Model) Predict(imageFile string) ([]float32, error) {
	imageNet := vision.NewImageNet()
	image, err := imageNet.LoadImageAndResize224(imageFile)
	if err != nil {
		log.Fatal(err)
	}
	input := image.MustUnsqueeze(0, true)
	raw_output := input.ApplyCModule(m.cmodule)
	output := raw_output.MustSoftmax(-1, gotch.Float, true)

	probs := output.Vals().([]float32)
	output.MustDrop()

	return probs, nil
}

func main() {
	modelPath := "./model.pt"
	imageFile := "./image.jpg"

	model, err := NewModel(modelPath)
	if err != nil {
		panic(err)
	}

	n := 1000
	for i := 0; i < n; i++ {
		prob, err := model.Predict(imageFile)
		if err != nil {
			panic(err)
		}

		if i%100 == 0 && i > 0 {
			log.Printf("Done... %2d: %v\n", i, prob)
		}
	}
}

[Joe Bloxsome]

@sugarme

Ah I see, that makes sense! Thanks for your help here as it was my mistake and not a bug with gotch.

Excellent work with gotch btw, I shall be making use of it in a few of my upcoming projects so do let me know if there's anywhere I can donate to support the project!