#include <execution_context.hpp>

Inheritance diagram for lbann::kfac_block< Device >:

Collaboration diagram for lbann::kfac_block< Device >:

Public Member Functions
	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

virtual void	on_forward_prop_end (lbann_comm *comm)

virtual int	get_local_memory_consumption ()=0
	Get local Memory Consumption. More...

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

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

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

virtual 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)
	Compute the inverse of the average Kronecker factors. More...

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

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

virtual void	start_communication_forward_end (lbann_comm *comm)=0

virtual void	end_communication_forward_end (lbann_comm *comm)=0

virtual void	start_communication_backward_end (lbann_comm *comm)=0

virtual void	end_communication_backward_end (lbann_comm *comm)=0

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

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

virtual int	get_inverse_matrices_size (lbann_comm *comm)=0
	Get inverse matrices size (offset). More...

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

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

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

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 ()

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

Protected Member Functions
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
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...

Private Attributes
kfac::KFACExecutionContext *	m_context
	The execution context that created this block. TODO: Use its own workspace and remove this pointer. More...

Detailed Description

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

A building block for K-FAC.

Definition at line 40 of file kfac/execution_context.hpp.

Constructor & Destructor Documentation

◆ kfac_block()

template<El::Device Device>

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
	)

inline

Constructor.

Definition at line 60 of file kfac_block.hpp.

◆ ~kfac_block()

template<El::Device Device>

virtual lbann::kfac_block< Device >::~kfac_block ( )

virtualdefault

Member Function Documentation

◆ compute_local_kronecker_factors()

template<El::Device Device>

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

virtual

Compute Kronecker factors.

Reimplemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ compute_preconditioned_gradients()

template<El::Device Device>

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

virtual

Compute the inverse of the average Kronecker factors.

Reimplemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ end_communication_backward_end()

template<El::Device Device>

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

pure virtual

Implemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ end_communication_forward_end()

template<El::Device Device>

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

pure virtual

Implemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ get_current_batch_size()

template<El::Device Device>

El::Int lbann::kfac_block< Device >::get_current_batch_size ( )

inline

Definition at line 192 of file kfac_block.hpp.

◆ get_gradient_wrt_weight_buffer()

template<El::Device Device>

DataType* lbann::kfac_block< Device >::get_gradient_wrt_weight_buffer ( int index )

inline

Definition at line 187 of file kfac_block.hpp.

◆ get_info()

template<El::Device Device>

virtual std::string lbann::kfac_block< Device >::get_info ( ) const

inlinevirtual

Get block's information in one line.

Reimplemented in lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

Definition at line 159 of file kfac_block.hpp.

◆ get_input_size()

template<El::Device Device>

El::Int lbann::kfac_block< Device >::get_input_size ( )

inline

Definition at line 194 of file kfac_block.hpp.

◆ get_internal_matrix_info()

template<El::Device Device>

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

virtual

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

Reimplemented in lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_gru< Device >, lbann::kfac_block_channelwise_fc< Device >, and lbann::kfac_block_bn< Device >.

◆ get_inverse_matrices()

template<El::Device Device>

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

pure virtual

Copy inverse matrices to output buffer.

Implemented in lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_gru< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ get_inverse_matrices_size()

template<El::Device Device>

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

pure virtual

Get inverse matrices size (offset).

Implemented in lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_gru< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ get_inverse_matrices_size_vector()

template<El::Device Device>

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

pure virtual

Get inverse matrices size vector.

Implemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ get_inverse_proc_rank()

template<El::Device Device>

size_t lbann::kfac_block< Device >::get_inverse_proc_rank ( ) const

inline

Definition at line 170 of file kfac_block.hpp.

◆ get_local_activation_buffer()

template<El::Device Device>

DataType* lbann::kfac_block< Device >::get_local_activation_buffer ( int index )

inline

Definition at line 172 of file kfac_block.hpp.

◆ get_local_error_buffer()

template<El::Device Device>

DataType* lbann::kfac_block< Device >::get_local_error_buffer ( int index )

inline

Definition at line 177 of file kfac_block.hpp.

◆ get_local_kronecker_buffers()

template<El::Device Device>

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

virtual

Get buffers of Kronecker factors for reduce-scatter.

Reimplemented in lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_gru< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ get_local_memory_consumption()

template<El::Device Device>

virtual int lbann::kfac_block< Device >::get_local_memory_consumption ( )

pure virtual

Get local Memory Consumption.

Implemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ get_name()

template<El::Device Device>

std::string lbann::kfac_block< Device >::get_name ( ) const

inline

Definition at line 168 of file kfac_block.hpp.

◆ get_output_size()

template<El::Device Device>

El::Int lbann::kfac_block< Device >::get_output_size ( )

inline

Definition at line 196 of file kfac_block.hpp.

Here is the call graph for this function:

◆ get_preconditioned_grad_buffers()

template<El::Device Device>

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

virtual

Get buffers of preconditioned parameter gradients.

Reimplemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ get_sync_info()

template<El::Device Device>

El::SyncInfo<Device> lbann::kfac_block< Device >::get_sync_info ( )

protected

Return the default sync info that may used in update functions.

◆ get_weight_buffer()

template<El::Device Device>

DataType* lbann::kfac_block< Device >::get_weight_buffer ( int index )

inline

Definition at line 182 of file kfac_block.hpp.

◆ get_workspace_matrix()

template<El::Device Device>

El::Matrix<DataType, Device>& lbann::kfac_block< Device >::get_workspace_matrix	(	const std::string &	key,
		size_t	height,
		size_t	width
	)

protected

Gets the Kronecker factor matrix of a FC layer. The same key is tied with the same matrix instance.

◆ initialize_activations_and_errors()

template<El::Device Device>

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

virtual

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

Reimplemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ on_forward_prop_end()

template<El::Device Device>

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

inlinevirtual

Reimplemented in lbann::kfac_block_gru< Device >.

Definition at line 81 of file kfac_block.hpp.

Here is the call graph for this function:

◆ resize_inverse_matrices_size()

template<El::Device Device>

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

pure virtual

Get inverse matrices size vector.

Implemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ set_current_batch_size()

template<El::Device Device>

void lbann::kfac_block< Device >::set_current_batch_size ( El::Int batch_size )

inline

Definition at line 156 of file kfac_block.hpp.

◆ set_inverse_matrices()

template<El::Device Device>

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

pure virtual

Copy inverse matrices from output buffer.

Implemented in lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_gru< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ start_communication_backward_end()

template<El::Device Device>

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

pure virtual

Implemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ start_communication_forward_end()

template<El::Device Device>

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

pure virtual

Implemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ update_kronecker_average()

template<El::Device Device>

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

virtual

Update the average Kronecker factors.

Reimplemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

◆ update_kronecker_inverse()

template<El::Device Device>

virtual void lbann::kfac_block< 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
	)

virtual

Compute the inverse of the average Kronecker factors.

Reimplemented in lbann::kfac_block_gru< Device >, lbann::kfac_block_fc_conv< Device >, lbann::kfac_block_bn< Device >, and lbann::kfac_block_channelwise_fc< Device >.

Member Data Documentation

◆ m_activations_copy

template<El::Device Device>

std::vector<std::unique_ptr<AbsDistMat> > lbann::kfac_block< Device >::m_activations_copy

protected

Definition at line 224 of file kfac_block.hpp.

◆ m_batch_size

template<El::Device Device>

int lbann::kfac_block< Device >::m_batch_size

protected

Definition at line 237 of file kfac_block.hpp.

◆ m_child_local_errors

template<El::Device Device>

std::vector<std::unique_ptr<AbsDistMat> > lbann::kfac_block< Device >::m_child_local_errors

protected

Definition at line 224 of file kfac_block.hpp.

◆ m_context

template<El::Device Device>

kfac::KFACExecutionContext* lbann::kfac_block< Device >::m_context

private

The execution context that created this block. TODO: Use its own workspace and remove this pointer.

Definition at line 251 of file kfac_block.hpp.

◆ m_enable_copy_activations

template<El::Device Device>

bool lbann::kfac_block< Device >::m_enable_copy_activations

protected

Enable copying of activations to enhance async communication.

Definition at line 243 of file kfac_block.hpp.

◆ m_enable_copy_errors

template<El::Device Device>

bool lbann::kfac_block< Device >::m_enable_copy_errors

protected

Enable copying of errors to enhance async communication.

Definition at line 240 of file kfac_block.hpp.

◆ m_errors_copy

template<El::Device Device>

std::vector<std::unique_ptr<AbsDistMat> > lbann::kfac_block< Device >::m_errors_copy

protected

Definition at line 224 of file kfac_block.hpp.

◆ m_has_kronecker_inverse

template<El::Device Device>

bool lbann::kfac_block< Device >::m_has_kronecker_inverse

protected

Whether this block already has an inverse history.

Definition at line 246 of file kfac_block.hpp.

◆ m_input_size

template<El::Device Device>

int lbann::kfac_block< Device >::m_input_size

protected

feature size and batch size (used in primary -> secondary grid communication)

Definition at line 237 of file kfac_block.hpp.

◆ m_inverse_proc_rank

template<El::Device Device>

const int lbann::kfac_block< Device >::m_inverse_proc_rank

protected

The process ID which perform inverse on Kronecker.

Definition at line 221 of file kfac_block.hpp.

◆ m_layer

template<El::Device Device>

Layer* lbann::kfac_block< Device >::m_layer

protected

The target layer.

Definition at line 214 of file kfac_block.hpp.

◆ m_layer_id

template<El::Device Device>

const size_t lbann::kfac_block< Device >::m_layer_id

protected

The layer ID in the model. TODO: Remove this.

Definition at line 218 of file kfac_block.hpp.

◆ m_output_size

template<El::Device Device>

int lbann::kfac_block< Device >::m_output_size

protected

Definition at line 237 of file kfac_block.hpp.

◆ m_parent_local_activations

template<El::Device Device>

std::vector<std::unique_ptr<AbsDistMat> > lbann::kfac_block< Device >::m_parent_local_activations

protected

distributed martices for activations and gradients.

Definition at line 224 of file kfac_block.hpp.

◆ m_requests_backward_end

template<El::Device Device>

std::vector<kfac::ReqT> lbann::kfac_block< Device >::m_requests_backward_end

protected

Definition at line 233 of file kfac_block.hpp.

◆ m_requests_forward_end

template<El::Device Device>

std::vector<kfac::ReqT> lbann::kfac_block< Device >::m_requests_forward_end

protected

Definition at line 233 of file kfac_block.hpp.

◆ m_subset_matrix

template<El::Device Device>

std::vector<std::unique_ptr<AbsDistMat> > lbann::kfac_block< Device >::m_subset_matrix

protected

Definition at line 224 of file kfac_block.hpp.

◆ m_weight_gradients

template<El::Device Device>

std::vector<std::unique_ptr<AbsDistMat> > lbann::kfac_block< Device >::m_weight_gradients

protected

Definition at line 224 of file kfac_block.hpp.

◆ m_weight_values

template<El::Device Device>

std::vector<std::unique_ptr<AbsDistMat> > lbann::kfac_block< Device >::m_weight_values

protected

Translatebetweengrid funciton has a basic implementation for STAR,STAR distributed matrices. Therefore, using local matrices for weights.

Definition at line 231 of file kfac_block.hpp.

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

Public Member Functions

Protected Member Functions

Protected Attributes

Private Attributes

Detailed Description

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

Constructor & Destructor Documentation

◆ kfac_block()

◆ ~kfac_block()

Member Function Documentation

◆ compute_local_kronecker_factors()

◆ compute_preconditioned_gradients()

◆ end_communication_backward_end()

◆ end_communication_forward_end()

◆ get_current_batch_size()

◆ get_gradient_wrt_weight_buffer()

◆ get_info()

◆ get_input_size()

◆ get_internal_matrix_info()

◆ get_inverse_matrices()

◆ get_inverse_matrices_size()

◆ get_inverse_matrices_size_vector()

◆ get_inverse_proc_rank()

◆ get_local_activation_buffer()

◆ get_local_error_buffer()

◆ get_local_kronecker_buffers()

◆ get_local_memory_consumption()

◆ get_name()

◆ get_output_size()

◆ get_preconditioned_grad_buffers()

◆ get_sync_info()

◆ get_weight_buffer()

◆ get_workspace_matrix()

◆ initialize_activations_and_errors()

◆ on_forward_prop_end()

◆ resize_inverse_matrices_size()

◆ set_current_batch_size()

◆ set_inverse_matrices()

◆ start_communication_backward_end()

◆ start_communication_forward_end()

◆ update_kronecker_average()

◆ update_kronecker_inverse()

Member Data Documentation

◆ m_activations_copy

◆ m_batch_size

◆ m_child_local_errors

◆ m_context

◆ m_enable_copy_activations

◆ m_enable_copy_errors

◆ m_errors_copy

◆ m_has_kronecker_inverse

◆ m_input_size

◆ m_inverse_proc_rank

◆ m_layer

◆ m_layer_id

◆ m_output_size

◆ m_parent_local_activations

◆ m_requests_backward_end

◆ m_requests_forward_end

◆ m_subset_matrix

◆ m_weight_gradients

◆ m_weight_values

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