Mirror/waifu2x-caffe
1
0
Fork 0
mirror of https://github.com/lltcggie/waifu2x-caffe.git synced 2025-06-26 13:42:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

GPU使用の場合はCUDAのメモリ割り当て関数を使うようにした

Browse Source
This commit is contained in:
lltcggie 2015-06-05 01:40:03 +09:00
parent d74a020e5b
commit ee1a7a3662
2 changed files with 156 additions and 86 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

84
common/waifu2x.cpp
View File

@ -8,6 +8,7 @@
#include <boost/filesystem.hpp>
#include <boost/algorithm/string.hpp>
#include <chrono>
#include <cuda_runtime.h>
#if defined(WIN32) || defined(WIN64)
#include <Windows.h>
@ -37,8 +38,31 @@ const int ConvertInverseMode = CV_YUV2RGB;
static std::once_flag waifu2x_once_flag;
static std::once_flag waifu2x_cudnn_once_flag;
#ifndef CUDA_CHECK_WAIFU2X
#define CUDA_CHECK_WAIFU2X(condition) \
do { \
cudaError_t error = condition; \
if(error != cudaSuccess) throw error; \
} while (0)
#endif
Waifu2x::Waifu2x() : is_inited(false)
#define CUDA_HOST_SAFE_FREE(ptr) \
do { \
if (ptr) { \
cudaFreeHost(ptr); \
ptr = nullptr; \
} \
} while (0)
#define SAFE_DELETE_WAIFU2X(ptr) \
do { \
if (ptr) { \
delete [] ptr; \
ptr = nullptr; \
} \
} while (0)
Waifu2x::Waifu2x() : is_inited(false), isCuda(false), block(nullptr), dummy_data(nullptr), out_block(nullptr)
{
}
@ -376,8 +400,6 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
net->layer_by_name("conv7_layer"));
assert(conv7_layer);
input_layer->set_batch_size(batch_size);
const int WidthNum = Width / output_size;
@ -388,9 +410,6 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
const int input_block_plane_size = block_size * block_size;
const int output_block_plane_size = crop_size * crop_size;
std::vector<float> block(input_block_plane_size * batch_size, 0.0f);
std::vector<float> dummy_data(block.size(), 0.0f);
// 画像は(消費メモリの都合上)output_size*output_sizeに分けて再構築する
for (int num = 0; num < BlockNum; num += batch_size)
{
@ -418,7 +437,7 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
// 画像を直列に変換
{
float *fptr = block.data() + (input_block_plane_size * n);
float *fptr = block + (input_block_plane_size * n);
const float *uptr = (const float *)someborderimg.data;
const auto Line = someborderimg.step1();
@ -436,7 +455,7 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
}
// ネットワークに画像を入力
input_layer->Reset(block.data(), dummy_data.data(), block.size());
input_layer->Reset(block, dummy_data, input_block_plane_size * batch_size);
// 計算
auto out = net->ForwardPrefilled(nullptr);
@ -452,7 +471,7 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
else
ptr = b->gpu_data();
caffe::caffe_copy(output_block_plane_size * processNum, ptr, block.data());
caffe::caffe_copy(output_block_plane_size * processNum, ptr, out_block);
for (int n = 0; n < processNum; n++)
{
@ -462,7 +481,7 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
const int w = wn * output_size;
const int h = hn * output_size;
const float *fptr = block.data() + (output_block_plane_size * n);
const float *fptr = out_block + (output_block_plane_size * n);
// 結果を入力画像にコピー(後に処理する部分とここで上書きする部分は被らないから、入力画像を上書きしても大丈夫)
for (int i = 0; i < crop_size; i++)
@ -489,6 +508,8 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
if (ScaleRatio <= 0.0)
return eWaifu2xError_InvalidParameter;
try
{
mode = Mode;
noise_level = NoiseLevel;
scale_ratio = ScaleRatio;
@ -541,9 +562,15 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
return eWaifu2xError_FailedOpenModelFile;
if (process == "cpu")
{
caffe::Caffe::set_mode(caffe::Caffe::CPU);
isCuda = false;
}
else
{
caffe::Caffe::set_mode(caffe::Caffe::GPU);
isCuda = true;
}
if (mode == "noise" || mode == "noise_scale" || mode == "auto_scale")
{
@ -573,7 +600,31 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
return ret;
}
const int input_block_plane_size = block_size * block_size;
const int output_block_plane_size = crop_size * crop_size;
if (isCuda)
{
CUDA_CHECK_WAIFU2X(cudaHostAlloc(&block, sizeof(float) * input_block_plane_size * batch_size, cudaHostAllocWriteCombined));
CUDA_CHECK_WAIFU2X(cudaHostAlloc(&dummy_data, sizeof(float) * input_block_plane_size * batch_size, cudaHostAllocWriteCombined));
CUDA_CHECK_WAIFU2X(cudaHostAlloc(&out_block, sizeof(float) * output_block_plane_size * batch_size, cudaHostAllocDefault));
}
else
{
block = new float[input_block_plane_size * batch_size];
dummy_data = new float[input_block_plane_size * batch_size];
out_block = new float[output_block_plane_size * batch_size];
}
for (size_t i = 0; i < input_block_plane_size * batch_size; i++)
dummy_data[i] = 0.0f;
is_inited = true;
}
catch (...)
{
return eWaifu2xError_InvalidParameter;
}
return eWaifu2xError_OK;
}
@ -583,6 +634,19 @@ void Waifu2x::destroy()
net_noise.reset();
net_scale.reset();
if (isCuda)
{
CUDA_HOST_SAFE_FREE(block);
CUDA_HOST_SAFE_FREE(dummy_data);
CUDA_HOST_SAFE_FREE(out_block);
}
else
{
SAFE_DELETE_WAIFU2X(block);
SAFE_DELETE_WAIFU2X(dummy_data);
SAFE_DELETE_WAIFU2X(out_block);
}
is_inited = false;
}

6
common/waifu2x.h
View File

@ -65,9 +65,15 @@ private:
std::string model_dir;
std::string process;
bool isCuda;
boost::shared_ptr<caffe::Net<float>> net_noise;
boost::shared_ptr<caffe::Net<float>> net_scale;
float *block;
float *dummy_data;
float *out_block;
private:
eWaifu2xError LoadImage(cv::Mat &float_image, const std::string &input_file);
eWaifu2xError CreateBrightnessImage(const cv::Mat &float_image, cv::Mat &im);