309 lines
5.8 KiB
C
309 lines
5.8 KiB
C
#ifndef TORCHDEFINE_H_
|
||
#define TORCHDEFINE_H_
|
||
|
||
#include <stdio.h>
|
||
#include "malloc.h"
|
||
|
||
#define LENGTH 256
|
||
#define MAX_GPU_COUNT 8
|
||
#define CLS_MODEL_BASE (10)
|
||
#define SEG_MODEL_BASE (10)
|
||
#define DET_MODEL_BASE (0)
|
||
#define PER_MODEL_MAX_COUNT (10)
|
||
|
||
|
||
enum DNN_MODEL_TYPE
|
||
{
|
||
DNN_MODEL_TYPE_CLS=0,
|
||
DNN_MODEL_TYPE_DET=1,
|
||
DNN_MODEL_TYPE_SEG=2,
|
||
DNN_MODEL_TYPE_DET_YOLO = 3,
|
||
DNN_MODEL_TYPE_SEG_FILTER = 4,
|
||
};
|
||
|
||
enum DNN_MODEL_FILE_TYPE
|
||
{
|
||
DNN_MODEL_CAFFE=0,
|
||
DNN_MODEL_TORCH=1,
|
||
DNN_MODEL_TENSORRT=2,
|
||
DNN_MODEL_ONNX=3,
|
||
};
|
||
|
||
|
||
#ifndef _TrainInfo_
|
||
#define _TrainInfo_
|
||
struct TrainInfo {
|
||
float loss;
|
||
float accuracy;//用户指定iou阈值，置信度阈值下的Precison * Recall
|
||
int record_index; //训练执行的轮数
|
||
bool is_detection; //是否是训练定位模型
|
||
float mAP;//coco 计算方式下的 map@0.5:0.95
|
||
float mAP50;//coco 计算方式下的 map@0.5
|
||
float mAP75;//coco 计算方式下的 map@0.75
|
||
float mIOU; // 检出定位的与目标的平均iou
|
||
float accuracy_bg;
|
||
float accuracy_fg;
|
||
float accuracy_objectness;
|
||
bool is_multilabel;
|
||
float accuracy_fg_all;
|
||
};
|
||
#endif
|
||
|
||
#ifndef _DetResult_
|
||
#define _DetResult_
|
||
struct DetResult
|
||
{
|
||
int image_id;
|
||
int label;
|
||
char name[16];
|
||
float score;
|
||
int left, top, height, width;
|
||
};
|
||
#endif
|
||
|
||
#ifndef _GpuDeviceInfo_
|
||
#define _GpuDeviceInfo_
|
||
struct GpuDeviceInfo
|
||
{
|
||
int gpu_count;
|
||
char gpu_name[MAX_GPU_COUNT][LENGTH];
|
||
int gpu_memory_size[MAX_GPU_COUNT];//单位MB
|
||
int gpu_sharedMemPerBlock[MAX_GPU_COUNT];//单位KB
|
||
int gpu_maxThreadsPerBlock[MAX_GPU_COUNT];
|
||
};
|
||
#endif
|
||
|
||
#ifndef _YYDLL_ERRORMSG_
|
||
#define _YYDLL_ERRORMSG_
|
||
struct YYErrorMsg
|
||
{
|
||
char strErrorMsg[1024];
|
||
};
|
||
#endif
|
||
|
||
#ifndef _DnnInferInfo_
|
||
#define _DnnInferInfo_
|
||
struct DnnInferInfo
|
||
{
|
||
char strModelFile[256];
|
||
int iBatchSize;
|
||
int iInferMode;
|
||
int iDeviceIndex;
|
||
};
|
||
#endif
|
||
|
||
#ifndef _Shape_
|
||
#define _Shape_
|
||
struct Shape
|
||
{
|
||
Shape() { N = 1; C = 1; H = 1; W = 1; }
|
||
Shape(unsigned int n, unsigned int c, unsigned int h, unsigned int w)
|
||
{
|
||
N = n; C = c; H = h; W = w;
|
||
}
|
||
unsigned int count()
|
||
{
|
||
return N * C * H * W;
|
||
}
|
||
unsigned int N;
|
||
unsigned int C;
|
||
unsigned int H;
|
||
unsigned int W;
|
||
};
|
||
#endif // !_Shape_
|
||
|
||
#ifndef _SegResult_
|
||
struct SegResult
|
||
{
|
||
unsigned char* pLabel;
|
||
float* pProb;
|
||
Shape inferShape;
|
||
int offsetX;
|
||
int offsetY;
|
||
double dScale;
|
||
};
|
||
#endif
|
||
|
||
#ifndef _Datum_
|
||
#define _Datum_
|
||
struct Datum
|
||
{
|
||
Datum()
|
||
{
|
||
data = NULL;
|
||
data_item_size = 4;
|
||
alloc = false;
|
||
}
|
||
|
||
Datum(int iDataItemSize)
|
||
{
|
||
data = NULL;
|
||
data_item_size = iDataItemSize;
|
||
alloc = false;
|
||
}
|
||
|
||
Datum(void* pData, Shape tshape, int itemSize)
|
||
{
|
||
data = pData;
|
||
this->shape = tshape;
|
||
data_item_size = itemSize;
|
||
alloc = false;
|
||
}
|
||
|
||
int GetDataItemSize()
|
||
{
|
||
return data_item_size;
|
||
}
|
||
void* Getdata()
|
||
{
|
||
if (data != NULL)
|
||
{
|
||
return data;
|
||
}
|
||
else
|
||
{
|
||
unsigned int size = shape.count();
|
||
data = malloc(size * data_item_size);
|
||
alloc = true;
|
||
return data;
|
||
}
|
||
}
|
||
|
||
void SetDataItemSize(int iDataItemSize)
|
||
{
|
||
if (data != NULL && iDataItemSize > data_item_size )
|
||
{
|
||
free(data);
|
||
data = NULL;
|
||
alloc = false;
|
||
}
|
||
data_item_size = iDataItemSize;
|
||
|
||
if (NULL == data)
|
||
{
|
||
unsigned int size = shape.count();
|
||
data = malloc(size * data_item_size);
|
||
alloc = true;
|
||
}
|
||
}
|
||
|
||
void Reshape(Shape s)
|
||
{
|
||
if (data != NULL && s.count() > shape.count())
|
||
{
|
||
free(data);
|
||
data = NULL;
|
||
alloc = false;
|
||
}
|
||
shape = s;
|
||
if (NULL == data)
|
||
{
|
||
unsigned int size = shape.count();
|
||
data = malloc(size * data_item_size);
|
||
alloc = true;
|
||
}
|
||
}
|
||
~Datum()
|
||
{
|
||
if (data != NULL)
|
||
{
|
||
if (alloc)
|
||
{
|
||
free(data);
|
||
data = NULL;
|
||
}
|
||
}
|
||
}
|
||
Shape shape;
|
||
void* data;
|
||
int data_item_size;
|
||
bool alloc;
|
||
};
|
||
#endif // !_Datum_
|
||
|
||
|
||
|
||
|
||
#ifndef _TrainingParameter_
|
||
#define _TrainingParameter_
|
||
struct TrainingParameter
|
||
{
|
||
//enhance relate
|
||
bool flip;
|
||
bool blur;
|
||
bool hsv;
|
||
bool transform;
|
||
bool crop;
|
||
//train information
|
||
int train_size;//100000
|
||
int kLogInterval;//100
|
||
int save_model_iterval;//1000
|
||
int batch_size;
|
||
//model relate need to write to file
|
||
//如果参数不相同将进行重新训练
|
||
int use_model_type;//0: small model; 1: medium model; 2: big model
|
||
int channel;
|
||
int height;
|
||
int width;
|
||
|
||
char data_file_root[LENGTH];
|
||
char val_data_file_root[LENGTH];
|
||
//添加容易出错的图像，每个batch保证一定比率添加到训练中
|
||
char harddata_file_root[LENGTH];
|
||
float input_scale;//1
|
||
|
||
float learn_rate;//0.01
|
||
|
||
//save relate
|
||
char pretrained_model_path[LENGTH];
|
||
char save_model_path[LENGTH];
|
||
|
||
int train_iter; //训练迭代轮数
|
||
bool save_interval_flag; //是否保存训练过程中文件
|
||
bool vflip; //是否使用上下翻转
|
||
float max_rotate_angle; //使用旋转增强时最大旋转角度(+-范围)
|
||
int loss_type; //训练loss_type
|
||
float image_scale; //图像训练前预缩放比例
|
||
int infer_height; //推理时模型输入尺寸高度
|
||
int infer_width; //推理时模型输入尺寸宽度
|
||
double label_smooth; //是否平滑类别标签，（即人为预估标记偏差）
|
||
bool in_memory_train; //是否在内存中训练
|
||
int model_type; //增加模型类别,0表示分类,1表示定位，2表示分割
|
||
bool with_prob; //推理时是否要输出概率
|
||
double crop_min_scale; //如果启用crop增强，最大的随机尺度
|
||
double crop_max_scale; //如果启用crop增强，最小的随机尺度
|
||
|
||
/*------------------ for detect --------------------------------------*/
|
||
char background_images_path[LENGTH];//加水印的背景图
|
||
bool is_create_anchor;//
|
||
float pos_overlap_low;//overlap大于这个值的被认为前景目标
|
||
float neg_overlap_up;//overlap小于这个值的被认为背景目标，中间区域认为忽略的前景目标
|
||
float gamma; //focal loss parameter
|
||
float balance; //focal loss parameter
|
||
};
|
||
#endif
|
||
|
||
#ifndef _InferenceParameter_
|
||
#define _InferenceParameter_
|
||
struct InferenceParameterEx
|
||
{
|
||
int use_model_type;
|
||
int batch_size;
|
||
int input_channel;
|
||
int input_width;
|
||
int input_height;
|
||
int label_num;
|
||
double input_scale;
|
||
int model_type; //增加模型类别,0表示分类,1表示定位，2表示分割
|
||
float mean_values[4];
|
||
float scale_value[4];
|
||
int input_count;
|
||
char input_name[3][64];
|
||
int output_count;
|
||
char output_name[3][64];
|
||
double image_scale;
|
||
};
|
||
#endif
|
||
|
||
#endif
|