113 lines
4.7 KiB
C++
113 lines
4.7 KiB
C++
|
// ======================================================================== //
|
||
|
|
// Copyright 2009-2019 Intel Corporation //
|
||
|
|
// //
|
||
|
|
// Licensed under the Apache License, Version 2.0 (the "License"); //
|
||
|
|
// you may not use this file except in compliance with the License. //
|
||
|
|
// You may obtain a copy of the License at //
|
||
|
|
// //
|
||
|
|
// http://www.apache.org/licenses/LICENSE-2.0 //
|
||
|
|
// //
|
||
|
|
// Unless required by applicable law or agreed to in writing, software //
|
||
|
|
// distributed under the License is distributed on an "AS IS" BASIS, //
|
||
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. //
|
||
|
|
// See the License for the specific language governing permissions and //
|
||
|
|
// limitations under the License. //
|
||
|
|
// ======================================================================== //
|
||
|
|
|
||
|
|
#include "common/tensor.h"
|
||
|
|
#include "image.h"
|
||
|
|
#include "node.h"
|
||
|
|
#include "input_reorder.h"
|
||
|
|
#include "output_reorder.h"
|
||
|
|
#include "transfer_function.h"
|
||
|
|
|
||
|
|
#pragma once
|
||
|
|
|
||
|
|
namespace oidn {
|
||
|
|
|
||
|
|
// Progress state
|
||
|
|
struct Progress
|
||
|
|
{
|
||
|
|
ProgressMonitorFunction func;
|
||
|
|
void* userPtr;
|
||
|
|
int taskCount;
|
||
|
|
};
|
||
|
|
|
||
|
|
class Executable
|
||
|
|
{
|
||
|
|
public:
|
||
|
|
virtual ~Executable() {}
|
||
|
|
virtual void execute(const Progress& progress, int taskIndex) = 0;
|
||
|
|
};
|
||
|
|
|
||
|
|
template<int K>
|
||
|
|
class Network : public Executable
|
||
|
|
{
|
||
|
|
public:
|
||
|
|
Network(const Ref<Device>& device, const std::map<std::string, Tensor>& weightMap);
|
||
|
|
|
||
|
|
void execute(const Progress& progress, int taskIndex) override;
|
||
|
|
|
||
|
|
std::shared_ptr<memory> allocTensor(const memory::dims& dims,
|
||
|
|
memory::format_tag format = memory::format_tag::any,
|
||
|
|
void* data = nullptr);
|
||
|
|
|
||
|
|
std::shared_ptr<memory> castTensor(const memory::dims& dims,
|
||
|
|
const std::shared_ptr<memory>& src,
|
||
|
|
size_t srcOffset = 0,
|
||
|
|
memory::format_tag format = memory::format_tag::any);
|
||
|
|
|
||
|
|
std::shared_ptr<memory> castTensor(const memory::dims& dims,
|
||
|
|
const std::shared_ptr<memory>& src,
|
||
|
|
const memory::dims& srcOffset);
|
||
|
|
|
||
|
|
void zeroTensor(const std::shared_ptr<memory>& dst);
|
||
|
|
|
||
|
|
memory::dims getInputReorderDims(const memory::dims& srcDims, int alignment);
|
||
|
|
|
||
|
|
std::shared_ptr<Node> addInputReorder(const Image& color,
|
||
|
|
const Image& albedo,
|
||
|
|
const Image& normal,
|
||
|
|
const std::shared_ptr<TransferFunction>& transferFunc,
|
||
|
|
int alignment,
|
||
|
|
const std::shared_ptr<memory>& userDst = nullptr);
|
||
|
|
|
||
|
|
std::shared_ptr<Node> addOutputReorder(const std::shared_ptr<memory>& src,
|
||
|
|
const std::shared_ptr<TransferFunction>& transferFunc,
|
||
|
|
const Image& output);
|
||
|
|
|
||
|
|
memory::dims getConvDims(const std::string& name, const memory::dims& srcDims);
|
||
|
|
std::shared_ptr<Node> addConv(const std::string& name,
|
||
|
|
const std::shared_ptr<memory>& src,
|
||
|
|
const std::shared_ptr<memory>& userDst = nullptr,
|
||
|
|
bool relu = true);
|
||
|
|
|
||
|
|
memory::dims getPoolDims(const memory::dims& srcDims);
|
||
|
|
std::shared_ptr<Node> addPool(const std::shared_ptr<memory>& src,
|
||
|
|
const std::shared_ptr<memory>& userDst = nullptr);
|
||
|
|
|
||
|
|
memory::dims getUpsampleDims(const memory::dims& srcDims);
|
||
|
|
std::shared_ptr<Node> addUpsample(const std::shared_ptr<memory>& src,
|
||
|
|
const std::shared_ptr<memory>& userDst = nullptr);
|
||
|
|
|
||
|
|
memory::dims getConcatDims(const memory::dims& src1Dims, const memory::dims& src2Dims);
|
||
|
|
|
||
|
|
std::shared_ptr<Node> addAutoexposure(const Image& color,
|
||
|
|
const std::shared_ptr<HDRTransferFunction>& transferFunc);
|
||
|
|
|
||
|
|
void finalize();
|
||
|
|
|
||
|
|
private:
|
||
|
|
Ref<Device> device;
|
||
|
|
engine eng;
|
||
|
|
stream sm;
|
||
|
|
std::vector<std::shared_ptr<Node>> nodes;
|
||
|
|
std::map<std::string, Tensor> weightMap;
|
||
|
|
|
||
|
|
// Memory allocation statistics
|
||
|
|
size_t activationAllocBytes = 0; // number of allocated activation bytes
|
||
|
|
size_t totalAllocBytes = 0; // total number of allocated bytes
|
||
|
|
};
|
||
|
|
|
||
|
|
} // namespace oidn
|