#include <kfac_block_channelwise_fc.hpp>

Inheritance diagram for lbann::kfac_block_channelwise_fc< Device >:

Collaboration diagram for lbann::kfac_block_channelwise_fc< Device >:

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

	kfac_block_channelwise_fc (const kfac_block_channelwise_fc &)=default

kfac_block_channelwise_fc &	operator= (const kfac_block_channelwise_fc &)=default

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

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

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

void	update_kronecker_average (lbann_comm *comm, DataType kronecker_decay, bool print_matrix, bool print_matrix_summary) final
	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) final
	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) final
	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) final
	Copies activations, errors, and weights from model class to private variables to be used in KFAC computation. More...

void	start_communication_forward_end (lbann_comm *comm) final

void	end_communication_forward_end (lbann_comm *comm) final

void	start_communication_backward_end (lbann_comm *comm) final

void	end_communication_backward_end (lbann_comm *comm) final

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

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

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

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

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

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

std::string	get_info () const final
	Get block's information in one line. 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

virtual void	on_forward_prop_end (lbann_comm *comm)

void	set_current_batch_size (El::Int batch_size)

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
convolution_layer< DataType, data_layout::DATA_PARALLEL, Device > *	get_conv_layer ()
	Get the pointer to its convolution_layer. More...

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...

Static Private Member Functions
static void	get_kronecker_factor_fc (El::AbstractMatrix< DataType > &factor, const El::AbstractMatrix< DataType > &activations, DataType alpha)
	Gets the Kronecker factor matrix of a FC layer. More...

static double	compute_pi (const El::Matrix< DataType, Device > &A, const El::Matrix< DataType, Device > &G, El::Matrix< DataType, Device > &ws, const El::SyncInfo< Device > &sync_info)
	Returns the pi constant. More...

Private Attributes
const bool	m_has_bias
	Information to perform its computation. More...

size_t	m_conv_input_spatial_prod

size_t	m_conv_output_spatial_prod

std::vector< int >	m_conv_input_spatial_dims

std::vector< int >	m_conv_output_spatial_dims

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

El::Matrix< DataType, Device >	m_kronecker_factor_buf_G

size_t	m_height_A
	The heights of the Kronecker factors. More...

size_t	m_height_G

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

El::Matrix< DataType, Device >	m_kronecker_average_G

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

El::Matrix< DataType, Device >	m_kronecker_inverse_G

size_t	m_Ainv_height = 0
	Size and height of inverse matrices. More...

size_t	m_Ainv_width = 0

size_t	m_Ginv_height = 0

size_t	m_Ginv_width = 0

El::Matrix< DataType, Device >	m_grad_buffer_v
	Vectorized gradient buffer (only for fully-connecter layers). More...

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_channelwise_fc< Device >

An FC/conv building block for K-FAC. TODO: Split into kfac_block_fc and kfac_block_conv.

Definition at line 38 of file kfac_block_channelwise_fc.hpp.

Constructor & Destructor Documentation

◆ kfac_block_channelwise_fc() [1/2]

template<El::Device Device>

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

inline

Constructor.

Definition at line 43 of file kfac_block_channelwise_fc.hpp.

Here is the call graph for this function:

◆ kfac_block_channelwise_fc() [2/2]

template<El::Device Device>

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

default

Member Function Documentation

◆ compute_local_kronecker_factors()

template<El::Device Device>

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

finalvirtual

Compute Kronecker factors.

Reimplemented from lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ compute_pi()

template<El::Device Device>

static double lbann::kfac_block_channelwise_fc< Device >::compute_pi	(	const El::Matrix< DataType, Device > &	A,
		const El::Matrix< DataType, Device > &	G,
		El::Matrix< DataType, Device > &	ws,
		const El::SyncInfo< Device > &	sync_info
	)

staticprivate

Returns the pi constant.

Here is the caller graph for this function:

◆ compute_preconditioned_gradients()

template<El::Device Device>

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

finalvirtual

Compute the inverse of the average Kronecker factors.

Reimplemented from lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ end_communication_backward_end()

template<El::Device Device>

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

finalvirtual

Implements lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ end_communication_forward_end()

template<El::Device Device>

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

finalvirtual

Implements lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ get_conv_layer()

template<El::Device Device>

convolution_layer<DataType, data_layout::DATA_PARALLEL, Device>* lbann::kfac_block_channelwise_fc< Device >::get_conv_layer ( )

inlineprivate

Get the pointer to its convolution_layer.

Definition at line 169 of file kfac_block_channelwise_fc.hpp.

Here is the call graph for this function:

◆ get_info()

template<El::Device Device>

std::string lbann::kfac_block_channelwise_fc< Device >::get_info ( ) const

inlinefinalvirtual

Get block's information in one line.

Reimplemented from lbann::kfac_block< Device >.

Definition at line 147 of file kfac_block_channelwise_fc.hpp.

Here is the call graph for this function:

◆ get_internal_matrix_info()

template<El::Device Device>

std::vector<std::tuple<std::string, size_t, size_t> > lbann::kfac_block_channelwise_fc< 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 >.

Here is the caller graph for this function:

◆ get_inverse_matrices()

template<El::Device Device>

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

finalvirtual

Copy inverse matrices to output buffer.

Implements lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ get_inverse_matrices_size()

template<El::Device Device>

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

finalvirtual

Get inverse matrices size (offset).

Implements lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ get_inverse_matrices_size_vector()

template<El::Device Device>

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

finalvirtual

Get inverse matrices size vector.

Implements lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ get_kronecker_factor_fc()

template<El::Device Device>

static void lbann::kfac_block_channelwise_fc< Device >::get_kronecker_factor_fc	(	El::AbstractMatrix< DataType > &	factor,
		const El::AbstractMatrix< DataType > &	activations,
		DataType	alpha
	)

staticprivate

Gets the Kronecker factor matrix of a FC layer.

Here is the caller graph for this function:

◆ get_local_kronecker_buffers()

template<El::Device Device>

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

inlinefinalvirtual

Get buffers of Kronecker factors for reduce-scatter.

Reimplemented from lbann::kfac_block< Device >.

Definition at line 90 of file kfac_block_channelwise_fc.hpp.

Here is the call graph for this function:

◆ get_local_memory_consumption()

template<El::Device Device>

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

inlinefinalvirtual

Get local Memory Consumption.

Implements lbann::kfac_block< Device >.

Definition at line 66 of file kfac_block_channelwise_fc.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_channelwise_fc< Device >::get_preconditioned_grad_buffers ( )

finalvirtual

Get buffers of preconditioned parameter gradients.

Reimplemented from lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ initialize_activations_and_errors()

template<El::Device Device>

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

finalvirtual

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

Reimplemented from lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ operator=()

template<El::Device Device>

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

default

Here is the caller graph for this function:

◆ resize_inverse_matrices_size()

template<El::Device Device>

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

finalvirtual

Get inverse matrices size vector.

Implements lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ set_inverse_matrices()

template<El::Device Device>

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

finalvirtual

Copy inverse matrices from output buffer.

Implements lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ start_communication_backward_end()

template<El::Device Device>

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

finalvirtual

Implements lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ start_communication_forward_end()

template<El::Device Device>

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

finalvirtual

Implements lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ update_kronecker_average()

template<El::Device Device>

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

finalvirtual

Update the average Kronecker factors.

Reimplemented from lbann::kfac_block< Device >.

Here is the caller graph for this function:

◆ update_kronecker_inverse()

template<El::Device Device>

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