File size: 3,482 Bytes
7873319 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | /*
* Copyright (c) 2021-2022, NVIDIA CORPORATION. All rights reserved.
*
* NVIDIA CORPORATION and its licensors retain all intellectual property
* and proprietary rights in and to this software, related documentation
* and any modifications thereto. Any use, reproduction, disclosure or
* distribution of this software and related documentation without an express
* license agreement from NVIDIA CORPORATION is strictly prohibited.
*/
/** @file trainable_buffer.cuh
* @author Thomas Müller, NVIDIA
* @brief An implementation of a trainable N-channel buffer within the tcnn API.
*/
#pragma once
#include <neural-graphics-primitives/common.h>
#include <neural-graphics-primitives/common_device.cuh>
#include <tiny-cuda-nn/common.h>
#include <tiny-cuda-nn/gpu_matrix.h>
#include <tiny-cuda-nn/gpu_memory.h>
#include <tiny-cuda-nn/network.h>
NGP_NAMESPACE_BEGIN
template <uint32_t N_DIMS, uint32_t RANK, typename T>
class TrainableBuffer : public tcnn::DifferentiableObject<float, T, T> {
using ResVector = Eigen::Matrix<int, RANK, 1>;
public:
TrainableBuffer(const ResVector& resolution) : m_resolution{resolution} {
m_params_gradient_weight.resize(n_params());
}
virtual ~TrainableBuffer() { }
void inference_mixed_precision_impl(cudaStream_t stream, const tcnn::GPUMatrixDynamic<float>& input, tcnn::GPUMatrixDynamic<T>& output, bool use_inference_matrices = true) override {
throw std::runtime_error{"The trainable buffer does not support inference(). Its content is meant to be used externally."};
}
std::unique_ptr<tcnn::Context> forward_impl(cudaStream_t stream, const tcnn::GPUMatrixDynamic<float>& input, tcnn::GPUMatrixDynamic<T>* output = nullptr, bool use_inference_matrices = false, bool prepare_input_gradients = false) override {
throw std::runtime_error{"The trainable buffer does not support forward(). Its content is meant to be used externally."};
}
void backward_impl(
cudaStream_t stream,
const tcnn::Context& ctx,
const tcnn::GPUMatrixDynamic<float>& input,
const tcnn::GPUMatrixDynamic<T>& output,
const tcnn::GPUMatrixDynamic<T>& dL_doutput,
tcnn::GPUMatrixDynamic<float>* dL_dinput = nullptr,
bool use_inference_matrices = false,
tcnn::EGradientMode param_gradients_mode = tcnn::EGradientMode::Overwrite
) override {
throw std::runtime_error{"The trainable buffer does not support backward(). Its content is meant to be used externally."};
}
void set_params_impl(T* params, T* inference_params, T* gradients) override { }
void initialize_params(tcnn::pcg32& rnd, float* params_full_precision, float scale = 1) override {
// Initialize the buffer to zero from the GPU
CUDA_CHECK_THROW(cudaMemset(params_full_precision, 0, n_params()*sizeof(float)));
}
size_t n_params() const override {
return m_resolution.prod() * N_DIMS;
}
uint32_t input_width() const override {
return RANK;
}
uint32_t padded_output_width() const override {
return N_DIMS;
}
uint32_t output_width() const override {
return N_DIMS;
}
uint32_t required_input_alignment() const override {
return 1; // No alignment required
}
std::vector<std::pair<uint32_t, uint32_t>> layer_sizes() const override {
return {};
}
T* gradient_weights() const {
return m_params_gradient_weight.data();
}
tcnn::json hyperparams() const override {
return {
{"otype", "TrainableBuffer"},
};
}
private:
ResVector m_resolution;
tcnn::GPUMemory<T> m_params_gradient_weight;
};
NGP_NAMESPACE_END
|