fixed run and onnx and uff support

Former-commit-id: 32b162a970d986c06d7d1abb459d5c1d04fdc045 [formerly 16421fb3ed67f59ed008915f362acaa8c85f6ef1] [formerly 92d5a72 [formerly 31e8790]]
Former-commit-id: c93d688c24bcb5d7dd029aa8b99d55bd69155877 [formerly e323b474c51a518df534a54fa2914f249d10194d]
Former-commit-id: 53f79d47710d97a4818e3f7a43fe9b2c1fcd738c

Author: jtang10

.gitignore

@@ -2,6 +2,7 @@ build/*
 networks/*
 *.o
 *.so
+_fixtures
 
 tmp
 tensorrt

.vscode/settings.json

@@ -44,6 +44,20 @@
         "utility": "cpp",
         "valarray": "cpp",
         "optional": "cpp",
-        "string_view": "cpp"
+        "string_view": "cpp",
+        "cstdarg": "cpp",
+        "atomic": "cpp",
+        "strstream": "cpp",
+        "bitset": "cpp",
+        "complex": "cpp",
+        "algorithm": "cpp",
+        "iterator": "cpp",
+        "map": "cpp",
+        "memory_resource": "cpp",
+        "set": "cpp",
+        "string": "cpp",
+        "iomanip": "cpp",
+        "cfenv": "cpp",
+        "cinttypes": "cpp"
     }
 }

cbits/predictor.hpp

@@ -24,7 +24,7 @@ extern "C"
   {
     TensorRT_CaffeFormat = 1,
     TensorRT_OnnxFormat = 2,
-    TensorRT_TensorFlowFormat = 3
+    TensorRT_UffFormat = 3,
   } TensorRT_ModelFormat;
 
   typedef enum TensorRT_DType
@@ -63,8 +63,9 @@ extern "C"
   X(TensorRT_Double, double)
 
   PredictorHandle
-  NewTensorRTPredictor(TensorRT_ModelFormat model_format, char *deploy_file,
-                       char *weights_file, TensorRT_DType model_datatype,
+  NewTensorRTPredictor(TensorRT_ModelFormat model_format, 
+                       char **model_files,
+                       TensorRT_DType model_datatype,
                        char **input_layer_names, int32_t num_input_layer_names,
                        char **output_layer_names, int32_t num_output_layer_names,
                        int32_t batch_size);

examples/_fixtures/networks/bvlc_googlenet.caffemodel

@@ -1,13 +1,13 @@
 package main
 
 import (
-	"bufio"
 	"context"
 	"fmt"
 	"image"
-	"os"
+	"io/ioutil"
 	"path/filepath"
 	"sort"
+	"strings"
 
 	"github.com/Unknwon/com"
 	"github.com/anthonynsimon/bild/imgio"
@@ -28,19 +28,19 @@ import (
 )
 
 var (
-	batchSize   = 1
-	model       = "resnet50"
-	shape       = []int{1, 3, 224, 224}
-	mean        = []float32{123.68, 116.779, 103.939}
-	scale       = []float32{1.0, 1.0, 1.0}
-	imgDir, _   = filepath.Abs("../_fixtures")
-	imgPath     = filepath.Join(imgDir, "platypus.jpg")
-	graph_url   = "http://s3.amazonaws.com/store.carml.org/models/caffe/resnet50/ResNet-50-deploy.prototxt"
-	weights_url = "http://s3.amazonaws.com/store.carml.org/models/caffe/resnet50/ResNet-50-model.caffemodel"
-	synset_url  = "http://s3.amazonaws.com/store.carml.org/synsets/imagenet/synset1001.txt"
+	batchSize  = 1
+	model      = "resnet50"
+	shape      = []int{1, 3, 224, 224}
+	mean       = []float32{123.68, 116.779, 103.939}
+	scale      = []float32{1.0, 1.0, 1.0}
+	baseDir, _ = filepath.Abs("../../_fixtures")
+	imgPath    = filepath.Join(baseDir, "platypus.jpg")
+	graphURL   = "http://s3.amazonaws.com/store.carml.org/models/caffe/resnet50/ResNet-50-deploy.prototxt"
+	weightsURL = "http://s3.amazonaws.com/store.carml.org/models/caffe/resnet50/ResNet-50-model.caffemodel"
+	synsetURL  = "http://s3.amazonaws.com/store.carml.org/synsets/imagenet/synset.txt"
 )
 
-// convert go Image to 1-dim array
+// convert go RGB Image to 1D normalized RGB array
 func cvtRGBImageToNCHW1DArray(src image.Image, mean []float32, scale []float32) ([]float32, error) {
 	if src == nil {
 		return nil, fmt.Errorf("src image nil")
@@ -49,41 +49,41 @@ func cvtRGBImageToNCHW1DArray(src image.Image, mean []float32, scale []float32)
 	in := src.Bounds()
 	height := in.Max.Y - in.Min.Y // image height
 	width := in.Max.X - in.Min.X  // image width
+	stride := width * height      // image size per channel
 
 	out := make([]float32, 3*height*width)
 	for y := 0; y < height; y++ {
 		for x := 0; x < width; x++ {
 			r, g, b, _ := src.At(x+in.Min.X, y+in.Min.Y).RGBA()
-			out[y*width+x] = (float32(b) - mean[2]) / scale[2]
-			out[width*height+y*width+x] = (float32(g) - mean[1]) / scale[1]
-			out[2*width*height+y*width+x] = (float32(r) - mean[0]) / scale[0]
+			out[0*stride+y*width+x] = (float32(r>>8) - mean[0]) / scale[0]
+			out[1*stride+y*width+x] = (float32(g>>8) - mean[1]) / scale[1]
+			out[2*stride+y*width+x] = (float32(b>>8) - mean[2]) / scale[2]
 		}
 	}
 
 	return out, nil
 }
 
 func main() {
-  defer tracer.Close()
-  
-	baseDir, _ := filepath.Abs("../tmp")
+	defer tracer.Close()
+
 	dir := filepath.Join(baseDir, model)
 	graph := filepath.Join(dir, model+".prototxt")
 	weights := filepath.Join(dir, model+".caffemodel")
 	synset := filepath.Join(dir, "synset.txt")
 
 	if !com.IsFile(graph) {
-		if _, _, err := downloadmanager.DownloadFile(graph_url, graph); err != nil {
+		if _, _, err := downloadmanager.DownloadFile(graphURL, graph); err != nil {
 			panic(err)
 		}
 	}
 	if !com.IsFile(weights) {
-		if _, _, err := downloadmanager.DownloadFile(weights_url, weights); err != nil {
+		if _, _, err := downloadmanager.DownloadFile(weightsURL, weights); err != nil {
 			panic(err)
 		}
 	}
 	if !com.IsFile(synset) {
-		if _, _, err := downloadmanager.DownloadFile(synset_url, synset); err != nil {
+		if _, _, err := downloadmanager.DownloadFile(synsetURL, synset); err != nil {
 			panic(err)
 		}
 	}
@@ -123,6 +123,10 @@ func main() {
 		Shape: shape,
 		Dtype: gotensor.Float32,
 	}
+	out := options.Node{
+		Key:   "prob",
+		Dtype: gotensor.Float32,
+	}
 
 	predictor, err := tensorrt.New(
 		ctx,
@@ -132,26 +136,21 @@ func main() {
 		options.Weights([]byte(weights)),
 		options.BatchSize(batchSize),
 		options.InputNodes([]options.Node{in}),
-		options.OutputNodes([]options.Node{
-			options.Node{
-				Key:   "prob",
-				Dtype: gotensor.Float32,
-			},
-		}),
+		options.OutputNodes([]options.Node{out}),
 	)
 	if err != nil {
 		panic(fmt.Sprintf("%v", err))
 	}
 	defer predictor.Close()
 
-	for ii:=0; ii < 3; ii++ {
-	err = predictor.Predict(ctx, input)
-	if err != nil {
-		panic(err)
+	for ii := 0; ii < 3; ii++ {
+		err = predictor.Predict(ctx, input)
+		if err != nil {
+			panic(err)
+		}
 	}
-}
 
-	enableCupti := true 
+	enableCupti := true
 	var cu *cupti.CUPTI
 	if enableCupti {
 		cu, err = cupti.New(cupti.Context(ctx))
@@ -174,64 +173,49 @@ func main() {
 		cu.Close()
 	}
 
-	if true {
-		profBuffer, err := predictor.ReadProfile()
-		if err != nil {
-			panic(err)
-		}
+	profBuffer, err := predictor.ReadProfile()
+	if err != nil {
+		panic(err)
+	}
 
-		t, err := ctimer.New(profBuffer)
-		if err != nil {
-			panic(err)
-		}
-		t.Publish(ctx, tracer.FRAMEWORK_TRACE)
+	t, err := ctimer.New(profBuffer)
+	if err != nil {
+		panic(err)
+	}
+	t.Publish(ctx, tracer.FRAMEWORK_TRACE)
 
-		outputs, err := predictor.ReadPredictionOutputs(ctx)
-		if err != nil {
-			panic(err)
-		}
+	outputs, err := predictor.ReadPredictionOutputs(ctx)
+	if err != nil {
+		panic(err)
+	}
 
-		output := outputs[0]
+	output := outputs[0]
+	labelsFileContent, err := ioutil.ReadFile(synset)
+	if err != nil {
+		panic(err)
+	}
+	labels := strings.Split(string(labelsFileContent), "\n")
 
-		var labels []string
-		f, err := os.Open(synset)
-		if err != nil {
-			panic(err)
-		}
-		defer f.Close()
-		scanner := bufio.NewScanner(f)
-		for scanner.Scan() {
-			line := scanner.Text()
-			labels = append(labels, line)
-		}
+	features := make([]dlframework.Features, batchSize)
+	featuresLen := len(output) / batchSize
 
-		features := make([]dlframework.Features, batchSize)
-		featuresLen := len(output) / batchSize
-
-		for ii := 0; ii < batchSize; ii++ {
-			rprobs := make([]*dlframework.Feature, featuresLen)
-			for jj := 0; jj < featuresLen; jj++ {
-				rprobs[jj] = feature.New(
-					feature.ClassificationIndex(int32(jj)),
-					feature.ClassificationLabel(labels[jj]),
-					feature.Probability(output[ii*featuresLen+jj]),
-				)
-			}
-			sort.Sort(dlframework.Features(rprobs))
-			features[ii] = rprobs
+	for ii := 0; ii < batchSize; ii++ {
+		rprobs := make([]*dlframework.Feature, featuresLen)
+		for jj := 0; jj < featuresLen; jj++ {
+			rprobs[jj] = feature.New(
+				feature.ClassificationIndex(int32(jj)),
+				feature.ClassificationLabel(labels[jj]),
+				feature.Probability(output[ii*featuresLen+jj]),
+			)
 		}
+		sort.Sort(dlframework.Features(rprobs))
+		features[ii] = rprobs
+	}
 
-		if true {
-			results := features[0]
-			for i := 0; i < 3; i++ {
-				prediction := results[i]
-				pp.Println(prediction.Probability)
-				pp.Println(prediction.GetClassification().GetIndex())
-				pp.Println(prediction.GetClassification().GetLabel())
-			}
-		} else {
-			_ = features
-		}
+	results := features[0]
+	for i := 0; i < 3; i++ {
+		prediction := results[i]
+		pp.Println(prediction.Probability, prediction.GetClassification().GetIndex(), prediction.GetClassification().GetLabel())
 	}
 }
 

examples/_fixtures/networks/bvlc_googlenet.caffemodel.2.tensorcache

@@ -1,15 +1,30 @@
 package tensorrt
 
+import (
+	"strings"
+)
+
 // ModelFormat ...
 type ModelFormat int
 
 const (
-	ModelFormatCaffe      ModelFormat = 1
-	ModelFormatOnnx       ModelFormat = 2
-	ModelFormatTensorFlow ModelFormat = 3
-	ModelFormatUnknown    ModelFormat = 999
+	ModelFormatCaffe   ModelFormat = 1
+	ModelFormatOnnx    ModelFormat = 2
+	ModelFormatUff     ModelFormat = 3
+	ModelFormatUnknown ModelFormat = 999
 )
 
-func ClassifyModelFormat(paths ...string) ModelFormat {
-	return ModelFormatCaffe
+func ClassifyModelFormat(path string) ModelFormat {
+	var format ModelFormat
+	if strings.HasSuffix(path, "prototxt") {
+		format = ModelFormatCaffe
+	} else if strings.HasSuffix(path, "onnx") {
+		format = ModelFormatOnnx
+	} else if strings.HasSuffix(path, "uff") {
+		format = ModelFormatUff
+	} else {
+		format = ModelFormatUnknown
+	}
+
+	return format
 }