使用するGPUデバイスの指定に対応 #29

bin/lang/english.json

@@ -90,5 +90,6 @@
 "MessageNoOverwrite":"Already file exists in the destination: %s",
 "MessageCudaOldDeviceError":"Your CUDA device too old.\r\nPlease use Compute Capability 2.0 or more devices",
 "OK":"OK",
-"Cancel":"Cancel"
+"Cancel":"Cancel",
+"IDC_STATIC_USE_GPU_NO":"Use GPU No"
 }

bin/lang/japanese.json

@@ -90,5 +90,6 @@
 "MessageNoOverwrite":"出力先に既にファイルが存在します: %s",
 "MessageCudaOldDeviceError":"CUDAデバイスが古いです。\r\nCompute Capability 2.0以上のデバイスを使用してください",
 "OK":"OK",
-"Cancel":"キャンセル"
+"Cancel":"キャンセル",
+"IDC_STATIC_USE_GPU_NO":"使用GPU No"
 }

common/waifu2x.cpp

@@ -172,6 +172,41 @@ namespace
 	};
 
 	IgnoreErrorCV g_IgnoreErrorCV;
+
+	class CudaDeviceSet
+	{
+	private:
+		int orgDevice;
+		bool mIsSet;
+
+	public:
+		CudaDeviceSet(const std::string &process, const int devno) : orgDevice(0), mIsSet(false)
+		{
+			if (process == "gpu" || process == "cudnn")
+			{
+				int count = 0;
+				if (cudaGetDeviceCount(&count) != CUDA_SUCCESS)
+					return;
+
+				if (devno >= count || count < 0)
+					return;
+
+				if (cudaGetDevice(&orgDevice) != CUDA_SUCCESS)
+					return;
+
+				if (cudaSetDevice(devno) != CUDA_SUCCESS)
+					return;
+
+				mIsSet = true;
+			}
+		}
+
+		~CudaDeviceSet()
+		{
+			if (mIsSet)
+				cudaSetDevice(orgDevice);
+		}
+	};
 }
 
 template<typename BufType>
@@ -323,7 +358,7 @@ static Waifu2x::eWaifu2xError writeProtoBinary(const ::google::protobuf::Message
 }
 
 
-Waifu2x::Waifu2x() : is_inited(false), isCuda(false), input_block(nullptr), dummy_data(nullptr), output_block(nullptr)
+Waifu2x::Waifu2x() : is_inited(false), isCuda(false), input_block(nullptr), dummy_data(nullptr), output_block(nullptr), gpu_no(0)
 {
 }
 
@@ -928,6 +963,8 @@ Waifu2x::eWaifu2xError Waifu2x::LoadParameterFromJson(boost::shared_ptr<caffe::N
 // ネットワークを使って画像を再構築する
 Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<float>> net, cv::Mat &im)
 {
+	CudaDeviceSet devset(process, gpu_no);
+
 	const auto Height = im.size().height;
 	const auto Width = im.size().width;
 	const auto Line = im.step1();
@@ -1134,7 +1171,7 @@ Waifu2x::eWaifu2xError Waifu2x::ReconstructImage(boost::shared_ptr<caffe::Net<fl
 Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &Mode, const int NoiseLevel,
 	const boost::optional<double> ScaleRatio, const boost::optional<int> ScaleWidth, const boost::optional<int> ScaleHeight,
 	const boost::filesystem::path &ModelDir, const std::string &Process,
-	const boost::optional<int> OutputQuality, const int OutputDepth, const bool UseTTA, const int CropSize, const int BatchSize)
+	const boost::optional<int> OutputQuality, const int OutputDepth, const bool UseTTA, const int CropSize, const int BatchSize, const int GPUNo)
 {
 	Waifu2x::eWaifu2xError ret;
 
@@ -1175,6 +1212,7 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
 
 		crop_size = CropSize;
 		batch_size = BatchSize;
+		gpu_no = GPUNo;
 
 		inner_padding = layer_num;
 		outer_padding = 1;
@@ -1185,17 +1223,6 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
 
 		output_block_size = crop_size + (inner_padding + outer_padding - layer_num) * 2;
 
-		std::call_once(waifu2x_once_flag, [argc, argv]()
-		{
-			assert(argc >= 1);
-
-			int tmpargc = 1;
-			char* tmpargvv[] = { argv[0] };
-			char** tmpargv = tmpargvv;
-			// glog“™‚Ì<E2809A>‰Šú‰»
-			caffe::GlobalInit(&tmpargc, &tmpargv);
-		});
-
 		const auto cuDNNCheckStartTime = std::chrono::system_clock::now();
 
 		if (process == "gpu")
@@ -1209,6 +1236,19 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
 
 		const auto cuDNNCheckEndTime = std::chrono::system_clock::now();
 
+		CudaDeviceSet devset(process, gpu_no);
+
+		std::call_once(waifu2x_once_flag, [argc, argv]()
+		{
+			assert(argc >= 1);
+
+			int tmpargc = 1;
+			char* tmpargvv[] = { argv[0] };
+			char** tmpargv = tmpargvv;
+			// glog“™‚Ì<E2809A>‰Šú‰»
+			caffe::GlobalInit(&tmpargc, &tmpargv);
+		});
+
 		boost::filesystem::path mode_dir_path(model_dir);
 		if (!mode_dir_path.is_absolute()) // model_dirが相対パスなら絶対パスに直す
 		{
@@ -1287,6 +1327,8 @@ Waifu2x::eWaifu2xError Waifu2x::init(int argc, char** argv, const std::string &M
 
 void Waifu2x::destroy()
 {
+	CudaDeviceSet devset(process, gpu_no);
+
 	net_noise.reset();
 	net_scale.reset();
 

common/waifu2x.h

@@ -117,6 +117,8 @@ private:
 	boost::optional<int> output_quality;
 	int output_depth;
 
+	int gpu_no;
+
 private:
 	static eWaifu2xError LoadMat(cv::Mat &float_image, const boost::filesystem::path &input_file);
 	static eWaifu2xError LoadMatBySTBI(cv::Mat &float_image, const std::vector<char> &img_data);
@@ -162,7 +164,7 @@ public:
 		const boost::optional<double> scale_ratio, const boost::optional<int> scale_width, const boost::optional<int> scale_height,
 		const boost::filesystem::path &model_dir, const std::string &process,
 		const boost::optional<int> output_quality = boost::optional<int>(), const int output_depth = 8, const bool use_tta = false,
-		const int crop_size = 128, const int batch_size = 1);
+		const int crop_size = 128, const int batch_size = 1, const int gpu_no = 0);
 
 	void destroy();
 

waifu2x-caffe-gui/MainDialog.cpp

@@ -36,68 +36,81 @@ const TCHAR * const MultiFileStr = TEXT("(Multi File)");
 const UINT_PTR nIDEventTimeLeft = 1000;
 
 
-// http://stackoverflow.com/questions/10167382/boostfilesystem-get-relative-path
+namespace
-boost::filesystem::path relativePath(const boost::filesystem::path &path, const boost::filesystem::path &relative_to)
 {
-	// create absolute paths
+	template<typename T>
-	boost::filesystem::path p = boost::filesystem::absolute(path);
+	static tstring to_tstring(T val)
-	boost::filesystem::path r = boost::filesystem::absolute(relative_to);
-
-	// if root paths are different, return absolute path
-	if (p.root_path() != r.root_path())
-		return p;
-
-	// initialize relative path
-	boost::filesystem::path result;
-
-	// find out where the two paths diverge
-	boost::filesystem::path::const_iterator itr_path = p.begin();
-	boost::filesystem::path::const_iterator itr_relative_to = r.begin();
-	while (*itr_path == *itr_relative_to && itr_path != p.end() && itr_relative_to != r.end())
 	{
-		++itr_path;
+#ifdef UNICODE
-		++itr_relative_to;
+		return std::to_wstring(val);
+#else
+		return std::to_string(val);
+#endif
 	}
 
-	// add "../" for each remaining token in relative_to
+	// http://stackoverflow.com/questions/10167382/boostfilesystem-get-relative-path
-	if (itr_relative_to != r.end())
+	boost::filesystem::path relativePath(const boost::filesystem::path &path, const boost::filesystem::path &relative_to)
 	{
-		++itr_relative_to;
+		// create absolute paths
-		while (itr_relative_to != r.end())
+		boost::filesystem::path p = boost::filesystem::absolute(path);
+		boost::filesystem::path r = boost::filesystem::absolute(relative_to);
+
+		// if root paths are different, return absolute path
+		if (p.root_path() != r.root_path())
+			return p;
+
+		// initialize relative path
+		boost::filesystem::path result;
+
+		// find out where the two paths diverge
+		boost::filesystem::path::const_iterator itr_path = p.begin();
+		boost::filesystem::path::const_iterator itr_relative_to = r.begin();
+		while (*itr_path == *itr_relative_to && itr_path != p.end() && itr_relative_to != r.end())
 		{
-			result /= "..";
+			++itr_path;
 			++itr_relative_to;
 		}
+
+		// add "../" for each remaining token in relative_to
+		if (itr_relative_to != r.end())
+		{
+			++itr_relative_to;
+			while (itr_relative_to != r.end())
+			{
+				result /= "..";
+				++itr_relative_to;
+			}
+		}
+
+		// add remaining path
+		while (itr_path != p.end())
+		{
+			result /= *itr_path;
+			++itr_path;
+		}
+
+		return result;
 	}
 
-	// add remaining path
+	std::vector<int> CommonDivisorList(const int N)
-	while (itr_path != p.end())
 	{
-		result /= *itr_path;
+		std::vector<int> list;
-		++itr_path;
+
+		const int sq = sqrt(N);
+		for (int i = 1; i <= sq; i++)
+		{
+			if (N % i == 0)
+				list.push_back(i);
+		}
+
+		const int sqs = list.size();
+		for (int i = 0; i < sqs; i++)
+			list.push_back(N / list[i]);
+
+		std::sort(list.begin(), list.end());
+
+		return list;
 	}
-
-	return result;
-}
-
-std::vector<int> CommonDivisorList(const int N)
-{
-	std::vector<int> list;
-
-	const int sq = sqrt(N);
-	for (int i = 1; i <= sq; i++)
-	{
-		if (N % i == 0)
-			list.push_back(i);
-	}
-
-	const int sqs = list.size();
-	for (int i = 0; i < sqs; i++)
-		list.push_back(N / list[i]);
-
-	std::sort(list.begin(), list.end());
-
-	return list;
 }
 
 
@@ -675,7 +688,7 @@ void DialogEvent::ProcessWaifu2x()
 	}
 
 	Waifu2x w;
-	ret = w.init(__argc, __argv, mode, noise_level, ScaleRatio, ScaleWidth, ScaleHeight, model_dir, process, output_quality, output_depth, use_tta, crop_size, batch_size);
+	ret = w.init(__argc, __argv, mode, noise_level, ScaleRatio, ScaleWidth, ScaleHeight, model_dir, process, output_quality, output_depth, use_tta, crop_size, batch_size, gpu_no);
 	if (ret != Waifu2x::eWaifu2xError_OK)
 		SendMessage(dh, WM_ON_WAIFU2X_ERROR, (WPARAM)&ret, 0);
 	else
@@ -924,6 +937,8 @@ void DialogEvent::SaveIni(const bool isSyncMember)
 
 	WritePrivateProfileString(TEXT("Setting"), TEXT("LastUseTTA"), to_tstring(use_tta ? 1 : 0).c_str(), getTString(SettingFilePath).c_str());
 
+	WritePrivateProfileString(TEXT("Setting"), TEXT("LastDeviceNo"), to_tstring(gpu_no).c_str(), getTString(SettingFilePath).c_str());
+
 	if (output_quality)
 		WritePrivateProfileString(TEXT("Setting"), TEXT("LastOutputQuality"), boost::lexical_cast<tstring>(*output_quality).c_str(), getTString(SettingFilePath).c_str());
 	else
@@ -1130,7 +1145,7 @@ UINT_PTR DialogEvent::OFNHookProcOut(HWND hdlg, UINT uiMsg, WPARAM wParam, LPARA
 
 DialogEvent::DialogEvent() : dh(nullptr), mode("noise_scale"), noise_level(1), scale_ratio(2.0), scale_width(0), scale_height(0), model_dir(TEXT("models/anime_style_art_rgb")),
 process("gpu"), outputExt(TEXT(".png")), inputFileExt(TEXT("png:jpg:jpeg:tif:tiff:bmp:tga")),
-use_tta(false), output_depth(8), crop_size(128), batch_size(1), isLastError(false), scaleType(eScaleTypeEnd),
+use_tta(false), output_depth(8), crop_size(128), batch_size(1), gpu_no(0), isLastError(false), scaleType(eScaleTypeEnd),
 TimeLeftThread(-1), TimeLeftGetTimeThread(0), isCommandLineStart(false), tAutoMode(TEXT("none")),
 isArgStartAuto(true), isArgStartSuccessFinish(true), isOutputNoOverwrite(false), isNotSaveParam(false)
 {}
@@ -1650,6 +1665,8 @@ void DialogEvent::Create(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData)
 
 		use_tta = GetPrivateProfileInt(TEXT("Setting"), TEXT("LastUseTTA"), 0, getTString(SettingFilePath).c_str()) != 0;
 
+		gpu_no = GetPrivateProfileInt(TEXT("Setting"), TEXT("LastDeviceNo"), 0, getTString(SettingFilePath).c_str());
+
 		output_quality.reset();
 		const int num = GetPrivateProfileInt(TEXT("Setting"), TEXT("LastOutputQuality"), -100, getTString(SettingFilePath).c_str());
 		if (num != -100)
@@ -1920,6 +1937,10 @@ void DialogEvent::Create(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData)
 				L"input batch size", false,
 				1, L"int", cmd);
 
+			TCLAP::ValueArg<int> cmdGPUNoFile(L"", L"gpu",
+				L"gpu device no", false,
+				0, L"int", cmd);
+
 			std::vector<int> cmdBoolConstraintV;
 			cmdBoolConstraintV.push_back(0);
 			cmdBoolConstraintV.push_back(1);
@@ -2134,6 +2155,13 @@ void DialogEvent::Create(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData)
 					isSetParam = true;
 				}
 
+				if (cmdGPUNoFile.isSet())
+				{
+					gpu_no = cmdGPUNoFile.getValue();
+
+					isSetParam = true;
+				}
+
 				if (cmdTTA.isSet())
 				{
 					SendMessage(GetDlgItem(dh, IDC_CHECK_TTA), BM_SETCHECK, cmdTTA.getValue() ? BST_CHECKED : BST_UNCHECKED, 0);
@@ -2844,6 +2872,7 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
 		bool isOutputNoOverwrite;
 		tstring tInputDirFix;
 		tstring tOutputDirFix;
+		int gpu_no;
 
 	private:
 		void AppSettingDialogEvent::SetWindowTextLang()
@@ -2872,6 +2901,8 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
 
 			SET_WINDOW_TEXT(IDC_CHECK_OUTPUT_NO_OVERWIRITE);
 
+			SET_WINDOW_TEXT(IDC_STATIC_USE_GPU_NO);
+
 #undef SET_WINDOW_TEXT
 		}
 
@@ -2926,6 +2957,8 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
 
 			SetWindowText(GetDlgItem(hWnd, IDC_EDIT_INPUT_DIR_FIX), tInputDirFix.c_str());
 			SetWindowText(GetDlgItem(hWnd, IDC_EDIT_OUTPUT_DIR_FIX), tOutputDirFix.c_str());
+
+			SetWindowText(GetDlgItem(hWnd, IDC_EDIT_USE_GPU_NO), to_tstring(gpu_no).c_str());
 		}
 
 		void Close(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData)
@@ -2935,6 +2968,8 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
 
 		void sync()
 		{
+			TCHAR *ptr = nullptr;
+
 			if (SendMessage(GetDlgItem(dh, IDC_RADIO_AUTO_START_ONE), BM_GETCHECK, 0, 0))
 				tAutoMode = TEXT("one");
 			else if (SendMessage(GetDlgItem(dh, IDC_RADIO_AUTO_START_MULTI), BM_GETCHECK, 0, 0))
@@ -2960,6 +2995,12 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
 			GetWindowText(GetDlgItem(dh, IDC_EDIT_OUTPUT_DIR_FIX), buf, _countof(buf));
 			buf[_countof(buf) - 1] = TEXT('\0');
 			tOutputDirFix = buf;
+
+			GetWindowText(GetDlgItem(dh, IDC_EDIT_USE_GPU_NO), buf, _countof(buf));
+			buf[_countof(buf) - 1] = TEXT('\0');
+			const auto devno = _tcstol(buf, &ptr, 10);
+			if ((!ptr || *ptr == TEXT('\0')) && devno >= 0)
+				gpu_no = devno;
 		}
 
 		void AppSettingDialogEvent::OK(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpData)
@@ -2983,6 +3024,7 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
 	cAppSettingDialogEvent.isOutputNoOverwrite = isOutputNoOverwrite;
 	cAppSettingDialogEvent.tInputDirFix = tInputDirFix;
 	cAppSettingDialogEvent.tOutputDirFix = tOutputDirFix;
+	cAppSettingDialogEvent.gpu_no = gpu_no;
 
 	CDialog cDialog;
 
@@ -3003,6 +3045,7 @@ void DialogEvent::AppSetting(HWND hWnd, WPARAM wParam, LPARAM lParam, LPVOID lpD
 		isOutputNoOverwrite = cAppSettingDialogEvent.isOutputNoOverwrite;
 		tInputDirFix = cAppSettingDialogEvent.tInputDirFix;
 		tOutputDirFix = cAppSettingDialogEvent.tOutputDirFix;
+		gpu_no = cAppSettingDialogEvent.gpu_no;
 
 		if (tOutputDirFix.length() > 0 && boost::filesystem::exists(tOutputDirFix))
 		{

waifu2x-caffe-gui/MainDialog.h

@@ -51,6 +51,8 @@ private:
 	int crop_size;
 	int batch_size;
 
+	int gpu_no;
+
 	std::vector<tstring> extList;
 
 	std::thread processThread;
@@ -107,16 +109,6 @@ private:
 	bool isNotSaveParam;
 
 private:
-	template<typename T>
-	static tstring to_tstring(T val)
-	{
-#ifdef UNICODE
-		return std::to_wstring(val);
-#else
-		return std::to_string(val);
-#endif
-	}
-
 	tstring AddName() const;
 
 	bool SyncMember(const bool NotSyncCropSize, const bool silent = false);

waifu2x-caffe/Source.cpp

@@ -135,6 +135,10 @@ int main(int argc, char** argv)
 		"input batch size", false,
 		1, "int", cmd);
 
+	TCLAP::ValueArg<int> cmdGPUNoFile("", "gpu",
+		"gpu device no", false,
+		0, "int", cmd);
+
 	std::vector<int> cmdTTAConstraintV;
 	cmdTTAConstraintV.push_back(0);
 	cmdTTAConstraintV.push_back(1);
@@ -319,7 +323,7 @@ int main(int argc, char** argv)
 	Waifu2x w;
 	ret = w.init(argc, argv, cmdMode.getValue(), cmdNRLevel.getValue(), ScaleRatio, ScaleWidth, ScaleHeight, cmdModelPath.getValue(), cmdProcess.getValue(),
 		cmdOutputQuality.getValue() == -1 ? boost::optional<int>() : cmdOutputQuality.getValue(), cmdOutputDepth.getValue(), use_tta, cmdCropSizeFile.getValue(),
-		cmdBatchSizeFile.getValue());
+		cmdBatchSizeFile.getValue(), cmdGPUNoFile.getValue());
 	switch (ret)
 	{
 	case Waifu2x::eWaifu2xError_InvalidParameter: