Skip to content

Commit

Permalink
added model downloader
Browse files Browse the repository at this point in the history
  • Loading branch information
@dusty-nv
dusty-nv committed Dec 17, 2022
1 parent 690dd4d commit a24401c
Show file tree
Hide file tree
Showing 6 changed files with 384 additions and 88 deletions.
70 changes: 43 additions & 27 deletions c/imageNet.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,23 +22,26 @@

#include "imageNet.h"
#include "tensorConvert.h"
#include "modelDownloader.h"

#include "cudaMappedMemory.h"
#include "cudaResize.h"

#include "commandLine.h"
#include "filesystem.h"
#include "logging.h"

#include <algorithm>



// constructor
imageNet::imageNet() : tensorNet()
{
mNetworkType = CUSTOM;
//mNetworkType = CUSTOM;
mThreshold = IMAGENET_DEFAULT_THRESHOLD;
mNumClasses = 0;
mThreshold = 0.01f;


mSmoothingBuffer = NULL;
mSmoothingFactor = 0;
}
Expand All @@ -50,37 +53,49 @@ imageNet::~imageNet()
CUDA_FREE_HOST(mSmoothingBuffer);
}


// Create
imageNet* imageNet::Create( imageNet::NetworkType networkType, uint32_t maxBatchSize,
imageNet* imageNet::Create( const char* network, uint32_t maxBatchSize,
precisionType precision, deviceType device, bool allowGPUFallback )
{
imageNet* net = new imageNet();
nlohmann::json model;

if( !net )
if( !DownloadModel(IMAGENET_MODEL_TYPE, network, model) )
return NULL;

if( !net->init(networkType, maxBatchSize, precision, device, allowGPUFallback) )
{
LogError(LOG_TRT "imageNet -- failed to initialize.\n");
return NULL;
}
std::string model_dir = "networks/" + model["dir"].get<std::string>() + "/";
std::string model_path = model_dir + model["model"].get<std::string>();
std::string prototxt = JSON_STR(model["prototxt"]);
std::string labels = JSON_STR(model["labels"]);
std::string input = JSON_STR_DEFAULT(model["input"], IMAGENET_DEFAULT_INPUT);
std::string output = JSON_STR_DEFAULT(model["output"], IMAGENET_DEFAULT_OUTPUT);

if( prototxt.length() > 0 )
prototxt = model_dir + prototxt;

net->mNetworkType = networkType;
return net;
if( locateFile(labels).length() == 0 )
labels = model_dir + labels;

return Create(prototxt.c_str(), model_path.c_str(), NULL,
labels.c_str(), input.c_str(), output.c_str(),
maxBatchSize, precision, device, allowGPUFallback);
}


// Create
imageNet* imageNet::Create( const char* prototxt_path, const char* model_path, const char* mean_binary,
const char* class_path, const char* input, const char* output, uint32_t maxBatchSize,
precisionType precision, deviceType device, bool allowGPUFallback )
{
// check for built-in model string
NetworkType type = NetworkTypeFromStr(model_path);

if( type != CUSTOM )
return Create(type, maxBatchSize, precision, device, allowGPUFallback);
if( FindModel(IMAGENET_MODEL_TYPE, model_path) )
{
return Create(model_path, maxBatchSize, precision, device, allowGPUFallback);
}
else if( fileExtension(model_path).length() == 0 )
{
LogError(LOG_TRT "couldn't find built-in classification model '%s'\n", model_path);
return NULL;
}

// load custom model
imageNet* net = new imageNet();
Expand All @@ -97,7 +112,7 @@ imageNet* imageNet::Create( const char* prototxt_path, const char* model_path, c
return net;
}


#if 0
// init
bool imageNet::init( imageNet::NetworkType networkType, uint32_t maxBatchSize,
precisionType precision, deviceType device, bool allowGPUFallback )
Expand Down Expand Up @@ -125,7 +140,7 @@ bool imageNet::init( imageNet::NetworkType networkType, uint32_t maxBatchSize,
else
return NULL;
}

#endif

// init
bool imageNet::init(const char* prototxt_path, const char* model_path, const char* mean_binary, const char* class_path,
Expand Down Expand Up @@ -172,6 +187,7 @@ bool imageNet::init(const char* prototxt_path, const char* model_path, const cha
}


#if 0
// NetworkTypeFromStr
imageNet::NetworkType imageNet::NetworkTypeFromStr( const char* modelName )
{
Expand Down Expand Up @@ -226,7 +242,7 @@ const char* imageNet::NetworkTypeToStr( imageNet::NetworkType network )

return "Custom";
}

#endif

// Create
imageNet* imageNet::Create( int argc, char** argv )
Expand All @@ -247,9 +263,9 @@ imageNet* imageNet::Create( const commandLine& cmdLine )
modelName = cmdLine.GetString("model", "googlenet");

// parse the network type
const imageNet::NetworkType type = NetworkTypeFromStr(modelName);
//const imageNet::NetworkType type = NetworkTypeFromStr(modelName);

if( type == imageNet::CUSTOM )
if( !FindModel("classification", modelName) )//type == imageNet::CUSTOM )
{
const char* prototxt = cmdLine.GetString("prototxt");
const char* labels = cmdLine.GetString("labels");
Expand All @@ -269,7 +285,7 @@ imageNet* imageNet::Create( const commandLine& cmdLine )
else
{
// create from pretrained model
net = imageNet::Create(type);
net = imageNet::Create(modelName);
}

if( !net )
Expand Down Expand Up @@ -322,7 +338,7 @@ bool imageNet::preProcess( void* image, uint32_t width, uint32_t height, imageFo

PROFILER_BEGIN(PROFILER_PREPROCESS);

if( mNetworkType == imageNet::INCEPTION_V4 )
if( mModelFile == "Inception-v4.caffemodel" ) //mNetworkType == imageNet::INCEPTION_V4 )
{
// downsample, convert to band-sequential RGB, and apply pixel normalization
if( CUDA_FAILED(cudaTensorNormRGB(image, format, width, height,
Expand Down
67 changes: 28 additions & 39 deletions c/imageNet.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,17 @@
*/
#define IMAGENET_DEFAULT_OUTPUT "prob"

/**
* Default value of the minimum confidence threshold for classification.
* @ingroup imageNet
*/
#define IMAGENET_DEFAULT_THRESHOLD 0.01f

/**
* The model type for imageNet in data/networks/models.json
* @ingroup imageNet
*/
#define IMAGENET_MODEL_TYPE "classification"

/**
* Standard command-line options able to be passed to imageNet::Create()
Expand All @@ -61,7 +72,7 @@
" --labels=LABELS path to text file containing the labels for each class\n" \
" --input-blob=INPUT name of the input layer (default is '" IMAGENET_DEFAULT_INPUT "')\n" \
" --output-blob=OUTPUT name of the output layer (default is '" IMAGENET_DEFAULT_OUTPUT "')\n" \
" --threshold=CONF minimum confidence threshold for classification (default is 0.01)\n" \
" --threshold=CONF minimum confidence threshold for classification (default is 0.01)\n" \
" --smoothing=WEIGHT weight between [0,1] or number of frames (disabled by default)\n" \
" --profile enable layer profiling in TensorRT\n\n"

Expand All @@ -72,44 +83,22 @@
*/
class imageNet : public tensorNet
{
public:
/**
* Network choice enumeration.
*/
enum NetworkType
{
CUSTOM, /**< Custom model provided by the user */
ALEXNET, /**< AlexNet trained on 1000-class ILSVRC12 */
GOOGLENET, /**< GoogleNet trained 1000-class ILSVRC12 */
GOOGLENET_12, /**< GoogleNet trained on 12-class subset of ImageNet ILSVRC12 from the tutorial */
RESNET_18, /**< ResNet-18 trained on 1000-class ILSVRC15 */
RESNET_50, /**< ResNet-50 trained on 1000-class ILSVRC15 */
RESNET_101, /**< ResNet-101 trained on 1000-class ILSVRC15 */
RESNET_152, /**< ResNet-50 trained on 1000-class ILSVRC15 */
VGG_16, /**< VGG-16 trained on 1000-class ILSVRC14 */
VGG_19, /**< VGG-19 trained on 1000-class ILSVRC14 */
INCEPTION_V4, /**< Inception-v4 trained on 1000-class ILSVRC12 */
};

public:
/**
* Parse a string to one of the built-in pretrained models.
* Valid names are "alexnet", "googlenet", "googlenet-12", or "googlenet_12", ect.
* @returns one of the imageNet::NetworkType enums, or imageNet::CUSTOM on invalid string.
*/
static NetworkType NetworkTypeFromStr( const char* model_name );

/**
* Convert a NetworkType enum to a string.
*/
static const char* NetworkTypeToStr( NetworkType network );

/**
* Load a new network instance
* Load one of the following pre-trained models:
*
* - alexnet, googlenet, googlenet-12,
* - resnet-18, resnet-50, resnet-101, resnet-152,
* - vgg-16, vgg-19, inception-v4
*
* These are all 1000-class models trained on ImageNet ILSVRC,
* except for googlenet-12 which is a 12-class subset of ILSVRC.
*/
static imageNet* Create( NetworkType networkType=GOOGLENET, uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE,
static imageNet* Create( const char* network="googlenet",
uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE,
precisionType precision=TYPE_FASTEST,
deviceType device=DEVICE_GPU, bool allowGPUFallback=true );

/**
* Load a new network instance
* @param prototxt_path File path to the deployable network prototxt
Expand Down Expand Up @@ -260,7 +249,7 @@ class imageNet : public tensorNet
* Retrieve the path to the file containing the class descriptions.
*/
inline const char* GetClassPath() const { return mClassPath.c_str(); }

#if 0
/**
* Retrieve the network type (alexnet or googlenet)
*/
Expand All @@ -270,7 +259,7 @@ class imageNet : public tensorNet
* Retrieve a string describing the network name.
*/
inline const char* GetNetworkName() const { return NetworkTypeToStr(mNetworkType); }

#endif
/**
* Return the confidence threshold used for classification.
*/
Expand Down Expand Up @@ -312,7 +301,7 @@ class imageNet : public tensorNet
protected:
imageNet();

bool init( NetworkType networkType, uint32_t maxBatchSize, precisionType precision, deviceType device, bool allowGPUFallback );
//bool init( NetworkType networkType, uint32_t maxBatchSize, precisionType precision, deviceType device, bool allowGPUFallback );
bool init(const char* prototxt_path, const char* model_path, const char* mean_binary, const char* class_path, const char* input, const char* output, uint32_t maxBatchSize, precisionType precision, deviceType device, bool allowGPUFallback );
bool loadClassInfo( const char* filename, int expectedClasses=-1 );

Expand All @@ -326,7 +315,7 @@ class imageNet : public tensorNet
std::vector<std::string> mClassDesc;

std::string mClassPath;
NetworkType mNetworkType;
//NetworkType mNetworkType;

float* mSmoothingBuffer;
float mSmoothingFactor;
Expand Down
Loading

0 comments on commit a24401c

Please sign in to comment.