#include <kfac_block_gru.hpp>

Inheritance diagram for lbann::kfac_block_gru< Device >:

Collaboration diagram for lbann::kfac_block_gru< Device >:

Public Member Functions
	kfac_block_gru (Layer layer, kfac::KFACExecutionContext context, size_t layer_id, size_t inverse_proc_rank, bool enable_copy_errors, bool enable_copy_activations, int input_size, int output_size)

	kfac_block_gru (const kfac_block_gru &)=default

kfac_block_gru &	operator= (const kfac_block_gru &)=default

void	on_forward_prop_end (lbann_comm *comm) override

const std::vector< El::AbstractMatrix< DataType > * >	get_local_kronecker_buffers () override
	Get buffers of Kronecker factors for reduce-scatter. More...

int	get_local_memory_consumption () override
	Get local Memory Consumption. More...

void	compute_local_kronecker_factors (lbann_comm *comm, bool print_matrix, bool print_matrix_summary) override
	Compute Kronecker factors. More...

void	update_kronecker_average (lbann_comm *comm, DataType kronecker_decay, bool print_matrix, bool print_matrix_summary) override
	Update the average Kronecker factors. More...

void	update_kronecker_inverse (lbann_comm *comm, bool use_pi, DataType damping_act, DataType damping_err, DataType learning_rate_factor, bool use_eigen_decomposition, bool print_matrix, bool print_matrix_summary, bool print_time) override
	Compute the inverse of the average Kronecker factors. More...

void	compute_preconditioned_gradients (lbann_comm *comm, DataType learning_rate_factor, bool print_matrix, bool print_matrix_summary, bool print_time) override
	Compute the inverse of the average Kronecker factors. More...

void	initialize_activations_and_errors (lbann_comm *comm, int num_local_activations, int num_local_errors, int num_weights) override
	Copies activations, errors, and weights from model class to private variables to be used in KFAC computation. More...

int	get_inverse_matrices (El::Matrix< DataType, Device > &output, int offset) override
	Copy inverse matrices to output buffer. More...

int	get_inverse_matrices_size (lbann_comm *comm) override
	Get inverse matrices size (offset). More...

int	set_inverse_matrices (El::Matrix< DataType, Device > &workspace, int offset, lbann_comm *comm) override
	Copy inverse matrices from output buffer. More...

void	start_communication_forward_end (lbann_comm *comm) override

void	end_communication_forward_end (lbann_comm *comm) override

void	start_communication_backward_end (lbann_comm *comm) override

void	end_communication_backward_end (lbann_comm *comm) override

std::vector< int >	get_inverse_matrices_size_vector (lbann_comm *comm) override
	Get inverse matrices size vector. More...

void	resize_inverse_matrices_size (El::Matrix< double, El::Device::CPU > &inverse_matrices_size, int block_number) override
	Get inverse matrices size vector. More...

const std::vector< El::AbstractMatrix< DataType > * >	get_preconditioned_grad_buffers () override
	Get buffers of preconditioned parameter gradients. More...

Public Member Functions inherited from lbann::kfac_block< Device >
	kfac_block (Layer layer, kfac::KFACExecutionContext context, size_t layer_id, size_t inverse_proc_rank, bool enable_copy_errors, bool enable_copy_activations, int input_size, int output_size)

virtual	~kfac_block ()=default

void	set_current_batch_size (El::Int batch_size)

virtual std::string	get_info () const
	Get block's information in one line. More...

std::string	get_name () const

size_t	get_inverse_proc_rank () const

DataType *	get_local_activation_buffer (int index)

DataType *	get_local_error_buffer (int index)

DataType *	get_weight_buffer (int index)

DataType *	get_gradient_wrt_weight_buffer (int index)

El::Int	get_current_batch_size ()

El::Int	get_input_size ()

El::Int	get_output_size ()

Private Member Functions
void	check_dnn_lib_spec () const

void	get_r_i (El::Matrix< DataType, Device > &r, El::Matrix< DataType, Device > &i, const El::Matrix< DataType, Device > &biases_ones, const El::Matrix< DataType, Device > &local_inputs, const El::Matrix< DataType, Device > &local_outputs, const El::Matrix< DataType, Device > &h0, size_t local_batch_size, const El::SyncInfo< Device > &sync_info)
	Recompute or copy (from cuDNN's reserve space if available) forward internal state (r and i). More...

void	get_weight_matrix (kfac_gru_util::weight_type matrix_type, El::Matrix< DataType, Device > &view)
	Get the view of a weight matrix or a bias vector or its gradients. More...

void	get_gradient_matrix (kfac_gru_util::weight_type matrix_type, El::Matrix< DataType, Device > &view)

void	get_gradient_buffer (kfac_gru_util::weight_type matrix_type, El::Matrix< DataType, Device > &view)

std::vector< std::tuple< std::string, size_t, size_t > >	get_internal_matrix_info () const override
	Return the list of internal matrices' (name, height, width) for debugging. All internal matrices should be ready when this function is called. More...

gru_layer< DataType, data_layout::DATA_PARALLEL, Device > *	get_gru_layer () const
	Get the pointer to its GRU_layer. More...

size_t	get_input_size () const

size_t	get_hidden_size () const

size_t	get_seq_length () const

void	send_recv_reserve_space (lbann_comm *comm)

Private Attributes
hydrogen::simple_buffer< El::byte, Device >	m_reserve_space_fwd
	A copy of the reserve space after forward passes. More...

El::Matrix< DataType, Device >	m_kronecker_factor_buf_A_h
	Lower triangle buffers of Kronecker factors. More...

El::Matrix< DataType, Device >	m_kronecker_factor_buf_A_x

std::unordered_map< kfac_gru_util::weight_type, El::Matrix< DataType, Device > >	m_kronecker_factor_buf_G

El::Matrix< DataType, Device >	m_kronecker_average_A_h
	Exponential moving average of Kronecker factors. More...

El::Matrix< DataType, Device >	m_kronecker_average_A_x

std::unordered_map< kfac_gru_util::weight_type, El::Matrix< DataType, Device > >	m_kronecker_average_G

El::Matrix< DataType, Device >	m_kronecker_inverse_A_h
	Inverse of the average Kronecker factors. More...

El::Matrix< DataType, Device >	m_kronecker_inverse_A_x

std::unordered_map< kfac_gru_util::weight_type, El::Matrix< DataType, Device > >	m_kronecker_inverse_G

El::Matrix< DataType, Device >	m_grad_buffer_A_h

El::Matrix< DataType, Device >	m_grad_buffer_A_x

std::unordered_map< kfac_gru_util::weight_type, El::Matrix< DataType, Device > >	m_grad_buffer_G

size_t	m_reserve_space_fwd_size = 0

std::vector< kfac::ReqT >	m_requests_workspace

Additional Inherited Members
Protected Member Functions inherited from lbann::kfac_block< Device >
El::Matrix< DataType, Device > &	get_workspace_matrix (const std::string &key, size_t height, size_t width)
	Gets the Kronecker factor matrix of a FC layer. The same key is tied with the same matrix instance. More...

El::SyncInfo< Device >	get_sync_info ()
	Return the default sync info that may used in update functions. More...

Protected Attributes inherited from lbann::kfac_block< Device >
Layer *	m_layer
	The target layer. More...

const size_t	m_layer_id
	The layer ID in the model. TODO: Remove this. More...

const int	m_inverse_proc_rank
	The process ID which perform inverse on Kronecker. More...

std::vector< std::unique_ptr< AbsDistMat > >	m_parent_local_activations
	distributed martices for activations and gradients. More...

std::vector< std::unique_ptr< AbsDistMat > >	m_child_local_errors

std::vector< std::unique_ptr< AbsDistMat > >	m_weight_gradients

std::vector< std::unique_ptr< AbsDistMat > >	m_subset_matrix

std::vector< std::unique_ptr< AbsDistMat > >	m_errors_copy

std::vector< std::unique_ptr< AbsDistMat > >	m_activations_copy

std::vector< std::unique_ptr< AbsDistMat > >	m_weight_values
	Translatebetweengrid funciton has a basic implementation for STAR,STAR distributed matrices. Therefore, using local matrices for weights. More...

std::vector< kfac::ReqT >	m_requests_forward_end

std::vector< kfac::ReqT >	m_requests_backward_end

int	m_input_size
	feature size and batch size (used in primary -> secondary grid communication) More...

int	m_output_size

int	m_batch_size

bool	m_enable_copy_errors
	Enable copying of errors to enhance async communication. More...

bool	m_enable_copy_activations
	Enable copying of activations to enhance async communication. More...

bool	m_has_kronecker_inverse
	Whether this block already has an inverse history. More...

Detailed Description

template<El::Device Device>
class lbann::kfac_block_gru< Device >

A BN building block for K-FAC.

Definition at line 115 of file kfac_block_gru.hpp.

Constructor & Destructor Documentation

◆ kfac_block_gru() [1/2]

template<El::Device Device>

lbann::kfac_block_gru< Device >::kfac_block_gru	(	Layer *	layer,
		kfac::KFACExecutionContext *	context,
		size_t	layer_id,
		size_t	inverse_proc_rank,
		bool	enable_copy_errors,
		bool	enable_copy_activations,
		int	input_size,
		int	output_size
	)

inline

Constructor.

Definition at line 120 of file kfac_block_gru.hpp.

Here is the call graph for this function:

◆ kfac_block_gru() [2/2]

template<El::Device Device>

lbann::kfac_block_gru< Device >::kfac_block_gru ( const kfac_block_gru< Device > & )

default

Member Function Documentation

◆ check_dnn_lib_spec()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::check_dnn_lib_spec ( ) const

private

◆ compute_local_kronecker_factors()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::compute_local_kronecker_factors	(	lbann_comm *	comm,
		bool	print_matrix,
		bool	print_matrix_summary
	)

overridevirtual

Compute Kronecker factors.

Reimplemented from lbann::kfac_block< Device >.

◆ compute_preconditioned_gradients()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::compute_preconditioned_gradients	(	lbann_comm *	comm,
		DataType	learning_rate_factor,
		bool	print_matrix,
		bool	print_matrix_summary,
		bool	print_time
	)

overridevirtual

Compute the inverse of the average Kronecker factors.

Reimplemented from lbann::kfac_block< Device >.

◆ end_communication_backward_end()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::end_communication_backward_end ( lbann_comm * comm )

overridevirtual

Implements lbann::kfac_block< Device >.

◆ end_communication_forward_end()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::end_communication_forward_end ( lbann_comm * comm )

overridevirtual

Implements lbann::kfac_block< Device >.

◆ get_gradient_buffer()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::get_gradient_buffer	(	kfac_gru_util::weight_type	matrix_type,
		El::Matrix< DataType, Device > &	view
	)

private

◆ get_gradient_matrix()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::get_gradient_matrix	(	kfac_gru_util::weight_type	matrix_type,
		El::Matrix< DataType, Device > &	view
	)

private

◆ get_gru_layer()

template<El::Device Device>

gru_layer<DataType, data_layout::DATA_PARALLEL, Device>* lbann::kfac_block_gru< Device >::get_gru_layer ( ) const

inlineprivate

Get the pointer to its GRU_layer.

Definition at line 255 of file kfac_block_gru.hpp.

◆ get_hidden_size()

template<El::Device Device>

size_t lbann::kfac_block_gru< Device >::get_hidden_size ( ) const

inlineprivate

Definition at line 266 of file kfac_block_gru.hpp.

◆ get_input_size()

template<El::Device Device>

size_t lbann::kfac_block_gru< Device >::get_input_size ( ) const

inlineprivate

Definition at line 261 of file kfac_block_gru.hpp.

◆ get_internal_matrix_info()

template<El::Device Device>

std::vector<std::tuple<std::string, size_t, size_t> > lbann::kfac_block_gru< Device >::get_internal_matrix_info ( ) const

overrideprivatevirtual

Return the list of internal matrices' (name, height, width) for debugging. All internal matrices should be ready when this function is called.

Reimplemented from lbann::kfac_block< Device >.

◆ get_inverse_matrices()

template<El::Device Device>

int lbann::kfac_block_gru< Device >::get_inverse_matrices	(	El::Matrix< DataType, Device > &	output,
		int	offset
	)

overridevirtual

Copy inverse matrices to output buffer.

Implements lbann::kfac_block< Device >.

◆ get_inverse_matrices_size()

template<El::Device Device>

int lbann::kfac_block_gru< Device >::get_inverse_matrices_size ( lbann_comm * comm )

overridevirtual

Get inverse matrices size (offset).

Implements lbann::kfac_block< Device >.

◆ get_inverse_matrices_size_vector()

template<El::Device Device>

std::vector<int> lbann::kfac_block_gru< Device >::get_inverse_matrices_size_vector ( lbann_comm * comm )

inlineoverridevirtual

Get inverse matrices size vector.

Implements lbann::kfac_block< Device >.

Definition at line 211 of file kfac_block_gru.hpp.

◆ get_local_kronecker_buffers()

template<El::Device Device>

const std::vector<El::AbstractMatrix<DataType>*> lbann::kfac_block_gru< Device >::get_local_kronecker_buffers ( )

overridevirtual

Get buffers of Kronecker factors for reduce-scatter.

Reimplemented from lbann::kfac_block< Device >.

◆ get_local_memory_consumption()

template<El::Device Device>

int lbann::kfac_block_gru< Device >::get_local_memory_consumption ( )

inlineoverridevirtual

Get local Memory Consumption.

Implements lbann::kfac_block< Device >.

Definition at line 156 of file kfac_block_gru.hpp.

Here is the call graph for this function:

◆ get_preconditioned_grad_buffers()

template<El::Device Device>

const std::vector<El::AbstractMatrix<DataType>*> lbann::kfac_block_gru< Device >::get_preconditioned_grad_buffers ( )

overridevirtual

Get buffers of preconditioned parameter gradients.

Reimplemented from lbann::kfac_block< Device >.

◆ get_r_i()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::get_r_i	(	El::Matrix< DataType, Device > &	r,
		El::Matrix< DataType, Device > &	i,
		const El::Matrix< DataType, Device > &	biases_ones,
		const El::Matrix< DataType, Device > &	local_inputs,
		const El::Matrix< DataType, Device > &	local_outputs,
		const El::Matrix< DataType, Device > &	h0,
		size_t	local_batch_size,
		const El::SyncInfo< Device > &	sync_info
	)

private

Recompute or copy (from cuDNN's reserve space if available) forward internal state (r and i).

◆ get_seq_length()

template<El::Device Device>

size_t lbann::kfac_block_gru< Device >::get_seq_length ( ) const

inlineprivate

Definition at line 267 of file kfac_block_gru.hpp.

◆ get_weight_matrix()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::get_weight_matrix	(	kfac_gru_util::weight_type	matrix_type,
		El::Matrix< DataType, Device > &	view
	)

private

Get the view of a weight matrix or a bias vector or its gradients.

◆ initialize_activations_and_errors()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::initialize_activations_and_errors	(	lbann_comm *	comm,
		int	num_local_activations,
		int	num_local_errors,
		int	num_weights
	)

overridevirtual

Copies activations, errors, and weights from model class to private variables to be used in KFAC computation.

Reimplemented from lbann::kfac_block< Device >.

◆ on_forward_prop_end()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::on_forward_prop_end ( lbann_comm * comm )

overridevirtual

Reimplemented from lbann::kfac_block< Device >.

◆ operator=()

template<El::Device Device>

kfac_block_gru& lbann::kfac_block_gru< Device >::operator= ( const kfac_block_gru< Device > & )

default

◆ resize_inverse_matrices_size()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::resize_inverse_matrices_size	(	El::Matrix< double, El::Device::CPU > &	inverse_matrices_size,
		int	block_number
	)

inlineoverridevirtual

Get inverse matrices size vector.

Implements lbann::kfac_block< Device >.

Definition at line 217 of file kfac_block_gru.hpp.

◆ send_recv_reserve_space()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::send_recv_reserve_space ( lbann_comm * comm )

private

◆ set_inverse_matrices()

template<El::Device Device>

int lbann::kfac_block_gru< Device >::set_inverse_matrices	(	El::Matrix< DataType, Device > &	workspace,
		int	offset,
		lbann_comm *	comm
	)

overridevirtual

Copy inverse matrices from output buffer.

Implements lbann::kfac_block< Device >.

◆ start_communication_backward_end()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::start_communication_backward_end ( lbann_comm * comm )

overridevirtual

Implements lbann::kfac_block< Device >.

◆ start_communication_forward_end()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::start_communication_forward_end ( lbann_comm * comm )

overridevirtual

Implements lbann::kfac_block< Device >.

◆ update_kronecker_average()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::update_kronecker_average	(	lbann_comm *	comm,
		DataType	kronecker_decay,
		bool	print_matrix,
		bool	print_matrix_summary
	)

overridevirtual

Update the average Kronecker factors.

Reimplemented from lbann::kfac_block< Device >.

◆ update_kronecker_inverse()

template<El::Device Device>

void lbann::kfac_block_gru< Device >::update_kronecker_inverse	(	lbann_comm *	comm,
		bool	use_pi,
		DataType	damping_act,
		DataType	damping_err,
		DataType	learning_rate_factor,
		bool	use_eigen_decomposition,
		bool	print_matrix,
		bool	print_matrix_summary,
		bool	print_time
	)

overridevirtual

Compute the inverse of the average Kronecker factors.

Reimplemented from lbann::kfac_block< Device >.

Member Data Documentation

◆ m_grad_buffer_A_h

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_gru< Device >::m_grad_buffer_A_h

private

Definition at line 294 of file kfac_block_gru.hpp.

◆ m_grad_buffer_A_x

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_gru< Device >::m_grad_buffer_A_x

private

Definition at line 294 of file kfac_block_gru.hpp.

◆ m_grad_buffer_G

template<El::Device Device>

std::unordered_map<kfac_gru_util::weight_type, El::Matrix<DataType, Device> > lbann::kfac_block_gru< Device >::m_grad_buffer_G

private

Definition at line 296 of file kfac_block_gru.hpp.

◆ m_kronecker_average_A_h

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_gru< Device >::m_kronecker_average_A_h

private

Exponential moving average of Kronecker factors.

Definition at line 285 of file kfac_block_gru.hpp.

◆ m_kronecker_average_A_x

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_gru< Device >::m_kronecker_average_A_x

private

Definition at line 285 of file kfac_block_gru.hpp.

◆ m_kronecker_average_G

template<El::Device Device>

std::unordered_map<kfac_gru_util::weight_type, El::Matrix<DataType, Device> > lbann::kfac_block_gru< Device >::m_kronecker_average_G

private

Definition at line 287 of file kfac_block_gru.hpp.

◆ m_kronecker_factor_buf_A_h

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_gru< Device >::m_kronecker_factor_buf_A_h

private

Lower triangle buffers of Kronecker factors.

Definition at line 279 of file kfac_block_gru.hpp.

◆ m_kronecker_factor_buf_A_x

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_gru< Device >::m_kronecker_factor_buf_A_x

private

Definition at line 279 of file kfac_block_gru.hpp.

◆ m_kronecker_factor_buf_G

template<El::Device Device>

std::unordered_map<kfac_gru_util::weight_type, El::Matrix<DataType, Device> > lbann::kfac_block_gru< Device >::m_kronecker_factor_buf_G

private

Definition at line 282 of file kfac_block_gru.hpp.

◆ m_kronecker_inverse_A_h

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_gru< Device >::m_kronecker_inverse_A_h

private

Inverse of the average Kronecker factors.

Definition at line 290 of file kfac_block_gru.hpp.

◆ m_kronecker_inverse_A_x

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_gru< Device >::m_kronecker_inverse_A_x

private

Definition at line 290 of file kfac_block_gru.hpp.

◆ m_kronecker_inverse_G

template<El::Device Device>

std::unordered_map<kfac_gru_util::weight_type, El::Matrix<DataType, Device> > lbann::kfac_block_gru< Device >::m_kronecker_inverse_G

private

Definition at line 292 of file kfac_block_gru.hpp.

◆ m_requests_workspace

template<El::Device Device>

std::vector<kfac::ReqT> lbann::kfac_block_gru< Device >::m_requests_workspace

private

Definition at line 300 of file kfac_block_gru.hpp.

◆ m_reserve_space_fwd

template<El::Device Device>

hydrogen::simple_buffer<El::byte, Device> lbann::kfac_block_gru< Device >::m_reserve_space_fwd

private

A copy of the reserve space after forward passes.

Definition at line 276 of file kfac_block_gru.hpp.

◆ m_reserve_space_fwd_size

template<El::Device Device>

size_t lbann::kfac_block_gru< Device >::m_reserve_space_fwd_size = 0

private

Definition at line 298 of file kfac_block_gru.hpp.

The documentation for this class was generated from the following file:

kfac_block_gru.hpp

Public Member Functions

Private Member Functions

Private Attributes

Additional Inherited Members

Detailed Description

template<El::Device Device> class lbann::kfac_block_gru< Device >

Constructor & Destructor Documentation

◆ kfac_block_gru() [1/2]

◆ kfac_block_gru() [2/2]

Member Function Documentation

◆ check_dnn_lib_spec()

◆ compute_local_kronecker_factors()

◆ compute_preconditioned_gradients()

◆ end_communication_backward_end()

◆ end_communication_forward_end()

◆ get_gradient_buffer()

◆ get_gradient_matrix()

◆ get_gru_layer()

◆ get_hidden_size()

◆ get_input_size()

◆ get_internal_matrix_info()

◆ get_inverse_matrices()

◆ get_inverse_matrices_size()

◆ get_inverse_matrices_size_vector()

◆ get_local_kronecker_buffers()

◆ get_local_memory_consumption()

◆ get_preconditioned_grad_buffers()

◆ get_r_i()

◆ get_seq_length()

◆ get_weight_matrix()

◆ initialize_activations_and_errors()

◆ on_forward_prop_end()

◆ operator=()

◆ resize_inverse_matrices_size()

◆ send_recv_reserve_space()

◆ set_inverse_matrices()

◆ start_communication_backward_end()

◆ start_communication_forward_end()

◆ update_kronecker_average()

◆ update_kronecker_inverse()

Member Data Documentation

◆ m_grad_buffer_A_h

◆ m_grad_buffer_A_x

◆ m_grad_buffer_G

◆ m_kronecker_average_A_h

◆ m_kronecker_average_A_x

◆ m_kronecker_average_G

◆ m_kronecker_factor_buf_A_h

◆ m_kronecker_factor_buf_A_x

◆ m_kronecker_factor_buf_G

◆ m_kronecker_inverse_A_h

◆ m_kronecker_inverse_A_x

◆ m_kronecker_inverse_G

◆ m_requests_workspace

◆ m_reserve_space_fwd

◆ m_reserve_space_fwd_size

template<El::Device Device>
class lbann::kfac_block_gru< Device >