C++ (Cpp) NetParameter::mutable_input_dim Examples

Example #1

File: waifu2x.cpp Project: nevercannot/waifu2x-caffe

Waifu2x::eWaifu2xError Waifu2x::SetParameter(caffe::NetParameter &param) const
{
	param.mutable_state()->set_phase(caffe::TEST);

	{
		auto mid = param.mutable_input_dim();

		if (mid->size() != 4)
			return eWaifu2xError_FailedParseModelFile;

		*mid->Mutable(0) = batch_size;
		*mid->Mutable(2) = input_block_size;
		*mid->Mutable(3) = input_block_size;
	}

	for (int i = 0; i < param.layer_size(); i++)
	{
		caffe::LayerParameter *layer_param = param.mutable_layer(i);
		const std::string& type = layer_param->type();
		if (type == "Convolution")
		{
			if (process == "cudnn")
				layer_param->mutable_convolution_param()->set_engine(caffe::ConvolutionParameter_Engine_CUDNN);
			else
				layer_param->mutable_convolution_param()->set_engine(caffe::ConvolutionParameter_Engine_CAFFE);
		}
		else if (type == "ReLU")
		{
			if (process == "cudnn")
				layer_param->mutable_relu_param()->set_engine(caffe::ReLUParameter_Engine_CUDNN);
			else
				layer_param->mutable_relu_param()->set_engine(caffe::ReLUParameter_Engine_CAFFE);
		}
	}

	return eWaifu2xError_OK;
}