Merge pull request #207 from NVlabs/overhaul-parameter-init

Overhaul parameter init

Author: Tom94

DOCUMENTATION.md

@@ -37,8 +37,6 @@ Lightning fast implementation of small multi-layer perceptrons (MLPs). Restricte
 	"n_neurons": 128,            // Neurons in each hidden layer.
 	                             // May only be 16, 32, 64, or 128.
 	"n_hidden_layers": 5,        // Number of hidden layers.
-	"feedback_alignment": false  // Use feedback alignment
-	                             // [Lillicrap et al. 2016].
 }
 ```
 

include/tiny-cuda-nn/cpp_api.h

@@ -86,7 +86,7 @@ class Module {
 		return m_param_precision;
 	}
 
-	virtual void initialize_params(size_t seed, float* params_full_precision) = 0;
+	virtual void initialize_params(size_t seed, float* params_full_precision, float scale = 1.0f) = 0;
 
 	virtual uint32_t n_output_dims() const = 0;
 	EPrecision output_precision() const {

include/tiny-cuda-nn/encoding.h

@@ -76,8 +76,8 @@ class Encoding : public DifferentiableObject<float, T, T> {
 	virtual MatrixLayout preferred_output_layout() const = 0;
 
 	// By default, an encoding has no parameters
-	void set_params(T* params, T* inference_params, T* backward_params, T* gradients) override { }
-	void initialize_params(pcg32& rnd, float* params_full_precision, T* params, T* inference_params, T* backward_params, T* gradients, float scale = 1) override { }
+	void set_params_impl(T* params, T* inference_params, T* gradients) override { }
+	void initialize_params(pcg32& rnd, float* params_full_precision, float scale = 1) override { }
 	size_t n_params() const override { return 0; }
 
 	std::vector<std::pair<uint32_t, uint32_t>> layer_sizes() const override { return {}; }

include/tiny-cuda-nn/encodings/composite.h

@@ -403,32 +403,18 @@ class CompositeEncoding : public Encoding<T> {
 		return m_nested.empty() ? AoS : m_nested.front()->preferred_output_layout();
 	}
 
-	void set_params(T* params, T* inference_params, T* backward_params, T* gradients) override {
+	void set_params_impl(T* params, T* inference_params, T* gradients) override {
 		size_t offset = 0;
 		for (auto& nested : m_nested) {
-			nested->set_params(
-				params + offset,
-				inference_params + offset,
-				backward_params + offset,
-				gradients + offset
-			);
+			nested->set_params(params + offset, inference_params + offset, gradients + offset);
 			offset += nested->n_params();
 		}
 	}
 
-	void initialize_params(pcg32& rnd, float* params_full_precision, T* params, T* inference_params, T* backward_params, T* gradients, float scale = 1) override {
-		size_t offset = 0;
+	void initialize_params(pcg32& rnd, float* params_full_precision, float scale = 1) override {
 		for (auto& nested : m_nested) {
-			nested->initialize_params(
-				rnd,
-				params_full_precision + offset,
-				params + offset,
-				inference_params + offset,
-				backward_params + offset,
-				gradients + offset,
-				scale
-			);
-			offset += nested->n_params();
+			nested->initialize_params(rnd, params_full_precision, scale);
+			params_full_precision += nested->n_params();
 		}
 	}
 

include/tiny-cuda-nn/encodings/grid.h

@@ -1079,7 +1079,7 @@ class GridEncodingTemplated : public GridEncoding<T> {
 			this->m_max_level_gpu,
 			m_interpolation_type,
 			m_grid_type,
-			use_inference_params ? m_grid_inference : m_grid,
+			use_inference_params ? this->inference_params() : this->params(),
 			forward->positions.data() ? forward->positions.view() : input.view(),
 			encoded_positions_soa,
 			forward->dy_dx.data()
@@ -1144,7 +1144,7 @@ class GridEncodingTemplated : public GridEncoding<T> {
 				grid_gradient_tmp = allocate_workspace(stream, m_n_params * sizeof(grad_t));
 				grid_gradient = (grad_t*)grid_gradient_tmp.data();
 			} else {
-				grid_gradient = (grad_t*)m_grid_gradient;
+				grid_gradient = (grad_t*)this->gradients();
 			}
 
 			if (param_gradients_mode == EGradientMode::Overwrite) {
@@ -1173,7 +1173,7 @@ class GridEncodingTemplated : public GridEncoding<T> {
 			);
 
 			if (!std::is_same<grad_t, T>::value) {
-				parallel_for_gpu(stream, n_params(), [grad=m_grid_gradient, grad_tmp=grid_gradient] __device__ (size_t i) {
+				parallel_for_gpu(stream, n_params(), [grad=this->gradients(), grad_tmp=grid_gradient] __device__ (size_t i) {
 					grad[i] = (T)grad_tmp[i];
 				});
 			}
@@ -1238,7 +1238,7 @@ class GridEncodingTemplated : public GridEncoding<T> {
 				grid_gradient_tmp = allocate_workspace(stream, m_n_params * sizeof(grad_t));
 				grid_gradient = (grad_t*)grid_gradient_tmp.data();
 			} else {
-				grid_gradient = (grad_t*)m_grid_gradient;
+				grid_gradient = (grad_t*)this->gradients();
 			}
 
 			if (param_gradients_mode == EGradientMode::Overwrite) {
@@ -1270,7 +1270,7 @@ class GridEncodingTemplated : public GridEncoding<T> {
 			);
 
 			if (!std::is_same<grad_t, T>::value) {
-				parallel_for_gpu(stream, n_params(), [grad=m_grid_gradient, grad_tmp=grid_gradient] __device__ (size_t i) {
+				parallel_for_gpu(stream, n_params(), [grad=this->gradients(), grad_tmp=grid_gradient] __device__ (size_t i) {
 					grad[i] = (T)grad_tmp[i];
 				});
 			}
@@ -1312,7 +1312,7 @@ class GridEncodingTemplated : public GridEncoding<T> {
 				dL_ddLdinput.view(),
 				forward.positions.data() ? forward.positions.view() : input.view(),
 				dL_dy_rm,
-				use_inference_params ? m_grid_inference : m_grid,
+				use_inference_params ? this->inference_params() : this->params(),
 				// outputs
 				dL_dinput->view()
 			);
@@ -1348,17 +1348,11 @@ class GridEncodingTemplated : public GridEncoding<T> {
 		return SoA;
 	}
 
-	void set_params(T* params, T* inference_params, T* backward_params, T* gradients) override {
-		m_grid = params;
-		m_grid_inference = inference_params;
-		m_grid_gradient = gradients;
-	}
-
-	void initialize_params(pcg32& rnd, float* params_full_precision, T* params, T* inference_params, T* backward_params, T* gradients, float scale = 1) override {
-		set_params(params, inference_params, backward_params, gradients);
+	void set_params_impl(T* params, T* inference_params, T* gradients) override { }
 
+	void initialize_params(pcg32& rnd, float* params_full_precision, float scale = 1) override {
 		// Initialize the hashgrid from the GPU, because the number of parameters can be quite large.
-		generate_random_uniform<float>(rnd, n_params(), params_full_precision, -1e-4f, 1e-4f);
+		generate_random_uniform<float>(rnd, n_params(), params_full_precision, -1e-4f * scale, 1e-4f * scale);
 	}
 
 	size_t n_params() const override {
@@ -1434,11 +1428,6 @@ class GridEncodingTemplated : public GridEncoding<T> {
 	bool m_stochastic_interpolation;
 	InterpolationType m_interpolation_type;
 	GridType m_grid_type;
-
-	// Storage of params
-	T* m_grid;
-	T* m_grid_inference;
-	T* m_grid_gradient;
 };
 
 template <typename T, uint32_t N_FEATURES_PER_LEVEL, HashType HASH_TYPE>

include/tiny-cuda-nn/network.h

@@ -34,12 +34,6 @@
 
 TCNN_NAMESPACE_BEGIN
 
-enum class WeightUsage {
-	Inference,
-	Forward,
-	Backward,
-};
-
 Activation string_to_activation(const std::string& activation_name);
 std::string to_string(Activation activation);
 

include/tiny-cuda-nn/network_with_input_encoding.h

@@ -110,48 +110,21 @@ class NetworkWithInputEncoding : public Network<float, T> {
 		}
 	}
 
-	void set_params(T* params, T* inference_params, T* backward_params, T* gradients) override {
+	void set_params_impl(T* params, T* inference_params, T* gradients) override {
 		size_t offset = 0;
-		m_network->set_params(
-			params + offset,
-			inference_params + offset,
-			backward_params + offset,
-			gradients + offset
-		);
+		m_network->set_params(params + offset, inference_params + offset, gradients + offset);
 		offset += m_network->n_params();
 
-		m_encoding->set_params(
-			params + offset,
-			inference_params + offset,
-			backward_params + offset,
-			gradients + offset
-		);
+		m_encoding->set_params(params + offset, inference_params + offset, gradients + offset);
 		offset += m_encoding->n_params();
 	}
 
-	void initialize_params(pcg32& rnd, float* params_full_precision, T* params, T* inference_params, T* backward_params, T* gradients, float scale = 1) override {
-		size_t offset = 0;
-		m_network->initialize_params(
-			rnd,
-			params_full_precision + offset,
-			params + offset,
-			inference_params + offset,
-			backward_params + offset,
-			gradients + offset,
-			scale
-		);
-		offset += m_network->n_params();
+	void initialize_params(pcg32& rnd, float* params_full_precision, float scale = 1) override {
+		m_network->initialize_params(rnd, params_full_precision, scale);
+		params_full_precision += m_network->n_params();
 
-		m_encoding->initialize_params(
-			rnd,
-			params_full_precision + offset,
-			params + offset,
-			inference_params + offset,
-			backward_params + offset,
-			gradients + offset,
-			scale
-		);
-		offset += m_encoding->n_params();
+		m_encoding->initialize_params(rnd, params_full_precision, scale);
+		params_full_precision += m_encoding->n_params();
 	}
 
 	size_t n_params() const override {

include/tiny-cuda-nn/networks/cutlass_mlp.h

@@ -66,8 +66,8 @@ class CutlassMLP : public Network<T> {
 		EGradientMode param_gradients_mode = EGradientMode::Overwrite
 	) override;
 
-	void set_params(T* params, T* inference_params, T* backward_params, T* gradients) override;
-	void initialize_params(pcg32& rnd, float* params_full_precision, T* params, T* inference_params, T* backward_params, T* gradients, float scale = 1) override;
+	void set_params_impl(T* params, T* inference_params, T* gradients) override;
+	void initialize_params(pcg32& rnd, float* params_full_precision, float scale = 1) override;
 
 	GPUMatrix<T, RM>& input_weight_matrix(bool inference) {
 		auto& weight_matrices = inference ? m_weight_matrices_inference : m_weight_matrices;

include/tiny-cuda-nn/networks/fully_fused_mlp.h

@@ -42,7 +42,7 @@ TCNN_NAMESPACE_BEGIN
 template <typename T, int WIDTH>
 class FullyFusedMLP : public Network<T> {
 public:
-	FullyFusedMLP(uint32_t input_width, uint32_t output_width, uint32_t n_hidden_layers, bool use_feedback_alignment, Activation activation, Activation output_activation);
+	FullyFusedMLP(uint32_t input_width, uint32_t output_width, uint32_t n_hidden_layers, Activation activation, Activation output_activation);
 
 	void inference_mixed_precision_impl(cudaStream_t stream, const GPUMatrixDynamic<T>& input, GPUMatrixDynamic<T>& output, bool use_inference_params = true) override;
 
@@ -59,37 +59,22 @@ class FullyFusedMLP : public Network<T> {
 		EGradientMode param_gradients_mode = EGradientMode::Overwrite
 	) override;
 
-	void set_params(T* params, T* inference_params, T* backward_params, T* gradients) override;
-	void initialize_params(pcg32& rnd, float* params_full_precision, T* params, T* inference_params, T* backward_params, T* gradients, float scale = 1) override;
+	void set_params_impl(T* params, T* inference_params, T* gradients) override;
+	void initialize_params(pcg32& rnd, float* params_full_precision, float scale = 1) override;
 
-	GPUMatrix<T, RM>& input_weight_matrix(WeightUsage usage) {
-		switch (usage) {
-			case WeightUsage::Inference: return m_weight_matrices_inference.front();
-			case WeightUsage::Forward: return m_weight_matrices.front();
-			case WeightUsage::Backward: return m_weight_matrices_backward.front();
-		}
-
-		throw std::runtime_error{"Invalid weight usage."};
+	GPUMatrix<T, RM>& input_weight_matrix(bool inference) {
+		auto& weight_matrices = inference ? m_weight_matrices_inference : m_weight_matrices;
+		return weight_matrices.front();
 	}
 
-	GPUMatrix<T, RM>& weight_matrix_at(WeightUsage usage, uint32_t idx) {
-		switch (usage) {
-			case WeightUsage::Inference: return m_weight_matrices_inference.at(1 + idx);
-			case WeightUsage::Forward: return m_weight_matrices.at(1 + idx);
-			case WeightUsage::Backward: return m_weight_matrices_backward.at(1 + idx);
-		}
-
-		throw std::runtime_error{"Invalid weight usage."};
+	GPUMatrix<T, RM>& weight_matrix_at(bool inference, uint32_t idx) {
+		auto& weight_matrices = inference ? m_weight_matrices_inference : m_weight_matrices;
+		return weight_matrices.at(1 + idx);
 	}
 
-	GPUMatrix<T, RM>& output_weight_matrix(WeightUsage usage) {
-		switch (usage) {
-			case WeightUsage::Inference: return m_weight_matrices_inference.back();
-			case WeightUsage::Forward: return m_weight_matrices.back();
-			case WeightUsage::Backward: return m_weight_matrices_backward.back();
-		}
-
-		throw std::runtime_error{"Invalid weight usage."};
+	GPUMatrix<T, RM>& output_weight_matrix(bool inference) {
+		auto& weight_matrices = inference ? m_weight_matrices_inference : m_weight_matrices;
+		return weight_matrices.back();
 	}
 
 	GPUMatrix<T, RM>& input_gradient_matrix() {
@@ -156,7 +141,6 @@ class FullyFusedMLP : public Network<T> {
 			{"output_activation", to_string(m_output_activation)},
 			{"n_neurons", m_network_width},
 			{"n_hidden_layers", m_n_hidden_layers},
-			{"feedback_alignment", m_use_feedback_alignment},
 		};
 	}
 
@@ -178,16 +162,11 @@ class FullyFusedMLP : public Network<T> {
 	Activation m_activation;
 	Activation m_output_activation;
 
-	bool m_use_feedback_alignment = false;
-
 	// Storage of params
 	std::vector<GPUMatrix<T, RM>> m_weight_matrices;
 	std::vector<GPUMatrix<T, RM>> m_weight_matrices_inference;
-	std::vector<GPUMatrix<T, RM>> m_weight_matrices_backward;
 	size_t m_total_n_params;
 
-	std::vector<GPUMatrix<float, RM>> m_weight_matrices_full_precision;
-
 	std::vector<GPUMatrix<T, RM>> m_gradient_matrices;
 };