Skip to content

Commit

Permalink
yolov5-2.0
Browse files Browse the repository at this point in the history
  • Loading branch information
enazoe committed Aug 31, 2020
1 parent 5be8ac0 commit 2151aab
Show file tree
Hide file tree
Showing 4 changed files with 15 additions and 11 deletions.
4 changes: 2 additions & 2 deletions configs/yolov5-3.0/yolov5s.cfg
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@

[net]
width=672
height=672
width=320
height=320
channels=3
nc=80
depth_multiple=0.33
Expand Down
9 changes: 7 additions & 2 deletions modules/hardswish.cu
Original file line number Diff line number Diff line change
Expand Up @@ -27,13 +27,15 @@ namespace nvinfer1
cudaDeviceProp prop;
cudaGetDeviceProperties(&prop, 0);
_n_max_thread_pre_block = prop.maxThreadsPerBlock;
// printf("Hardswish():%d\n", _n_max_thread_pre_block);
}

Hardswish::Hardswish(const void* data, size_t length)
{
const char *d = reinterpret_cast<const char*>(data), *a = d;
r(d, _n_max_thread_pre_block);
r(d, _n_output_size);
// printf("r:threads:%d,size:%d\n", _n_max_thread_pre_block, _n_output_size);
assert(d == a + length);
}

Expand Down Expand Up @@ -66,6 +68,7 @@ namespace nvinfer1
cudaStream_t stream_)
{
int n_data_size = n_batch_size_ * n_output_size_;
// printf("cuda_hardswish_layer:%d,size:%d\n", n_batch_size_, n_output_size_);
kernel_hardswish << <(n_data_size + threads_ -1)/threads_, threads_ >> >(
reinterpret_cast<const float*>(input_),
reinterpret_cast<float*>(output_),
Expand All @@ -76,7 +79,7 @@ namespace nvinfer1
int Hardswish::enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace,
cudaStream_t stream)
{
//printf("batch_size:%d,output_size:%d,threads:%d\n", batchSize, _n_output_size, _n_max_thread_pre_block);
// printf("batch_size:%d,output_size:%d,threads:%d\n", batchSize, _n_output_size, _n_max_thread_pre_block);
NV_CUDA_CHECK(cuda_hardswish_layer(inputs[0], outputs[0], batchSize, _n_output_size , _n_max_thread_pre_block,stream));
return 0;
}
Expand All @@ -91,21 +94,23 @@ namespace nvinfer1
char *d = static_cast<char*>(buffer), *a = d;
w(d, _n_max_thread_pre_block);
w(d, _n_output_size);
// printf("serialize:%d,%d\n", _n_max_thread_pre_block, _n_output_size);
assert(d == a + getSerializationSize());
}

void Hardswish::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput)
{

_n_output_size = in->dims.d[0] * in->dims.d[1] * in->dims.d[2];
// printf("output_size:%d,threads:%d\n", _n_output_size, _n_max_thread_pre_block);
// printf("configurePlugin:%d,%d,%d\n", in->dims.d[0], in->dims.d[1], in->dims.d[2]);
}
IPluginV2IOExt* Hardswish::clone() const
{
Hardswish *p = new Hardswish();
p->setPluginNamespace(_s_plugin_namespace.c_str());
p->_n_max_thread_pre_block = _n_max_thread_pre_block;
p->_n_output_size = _n_output_size;
// printf("clone:%d,%d\n", _n_max_thread_pre_block, _n_output_size);
return p;
}

Expand Down
4 changes: 2 additions & 2 deletions modules/trt_utils.h
Original file line number Diff line number Diff line change
Expand Up @@ -215,11 +215,11 @@ nvinfer1::ILayer * layer_conv_bn_act(
const int group_ =1,
const bool b_padding_ = true,
const bool b_bn_ = true,
const std::string s_act_ = "hardswish");
const std::string s_act_ = "leaky");

nvinfer1::ILayer * layer_act(nvinfer1::ITensor* input_,
nvinfer1::INetworkDefinition* network_,
const std::string s_act_ = "hardswish");
const std::string s_act_ = "leaky");

nvinfer1::ILayer * layer_bottleneck_csp(
std::string s_model_name_,
Expand Down
9 changes: 4 additions & 5 deletions samples/sample_detector.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,10 @@ int main()
Config config_v5;
config_v5.net_type = YOLOV5;
config_v5.detect_thresh = 0.5;
config_v5.file_model_cfg = "../configs/yolov5-3.0/yolov5s.cfg";
config_v5.file_model_weights = "../configs/yolov5-3.0/yolov5s.weights";
config_v5.calibration_image_list_file_txt = "../configs/calibration_images.txt";
config_v5.file_model_cfg = "../configs/yolov5-2.0/yolov5s.cfg";
config_v5.file_model_weights = "../configs/yolov5-2.0/yolov5s.weights";
config_v5.inference_precison = FP32;
config_v5.n_max_batch = 1;
//config_v5.n_max_batch = 2;

cv::Mat image0 = cv::imread("../configs/dog.jpg", cv::IMREAD_UNCHANGED);
cv::Mat image1 = cv::imread("../configs/person.jpg", cv::IMREAD_UNCHANGED);
Expand All @@ -58,7 +57,7 @@ int main()
cv::Mat temp0 = image0.clone();
cv::Mat temp1 = image1.clone();
batch_img.push_back(temp0);
//batch_img.push_back(temp1);
batch_img.push_back(temp1);

//detect
timer.reset();
Expand Down

0 comments on commit 2151aab

Please sign in to comment.