19 lines
691 B
C++
19 lines
691 B
C++
#ifndef TENSORRTX_LAYERS_H
|
|
#define TENSORRTX_LAYERS_H
|
|
|
|
#include <map>
|
|
#include <math.h>
|
|
#include <assert.h>
|
|
|
|
#include "NvInfer.h"
|
|
#include "cuda_runtime_api.h"
|
|
using namespace nvinfer1;
|
|
|
|
IScaleLayer *addBatchNorm2d(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, std::string lname, float eps);
|
|
|
|
IActivationLayer *bottleneck(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, int ch, int stride, std::string lname, int branch_type);
|
|
|
|
IActivationLayer *addConvRelu(INetworkDefinition *network, std::map<std::string, Weights> &weightMap, ITensor &input, int outch, int kernel, int stride, std::string lname);
|
|
|
|
#endif
|