#include <kfac_block_fc_conv.hpp>

Inheritance diagram for lbann::kfac_block_fc_conv< Device >:

Collaboration diagram for lbann::kfac_block_fc_conv< Device >:

Public Member Functions
	kfac_block_fc_conv (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, const bool is_conv)

	kfac_block_fc_conv (const kfac_block_fc_conv &)=default

kfac_block_fc_conv &	operator= (const kfac_block_fc_conv &)=default

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

const std::vector< El::AbstractMatrix< DataType > * >	get_local_kronecker_buffers () override
	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) 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...

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

const std::vector< El::AbstractMatrix< DataType > * >	get_preconditioned_grad_buffers () override
	Get buffers of preconditioned parameter gradients. 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...

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

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

std::string	get_info () const override
	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 void	get_kronecker_factor_conv (El::Matrix< DataType, Device > &factor, El::Matrix< DataType, Device > &Acol, const El::Matrix< DataType, Device > &activations, DataType alpha, size_t local_batch_size, size_t num_channels, const std::vector< int > &spatial_dims, const convolution_layer< DataType, data_layout::DATA_PARALLEL, Device > *l_conv, bool use_im2col, const El::SyncInfo< Device > &sync_info)
	Gets the Kronecker factor matrix of a convolutional 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_is_conv
	Information to perform its computation. More...

const bool	m_has_bias

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

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

Definition at line 71 of file kfac_block_fc_conv.hpp.

Constructor & Destructor Documentation

◆ kfac_block_fc_conv() [1/2]

template<El::Device Device>

lbann::kfac_block_fc_conv< Device >::kfac_block_fc_conv	(	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,
		const bool	is_conv
	)

inline

Constructor.

Definition at line 76 of file kfac_block_fc_conv.hpp.

Here is the call graph for this function:

◆ kfac_block_fc_conv() [2/2]

template<El::Device Device>

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

default

Member Function Documentation

◆ compute_local_kronecker_factors()

template<El::Device Device>

void lbann::kfac_block_fc_conv< 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_pi()

template<El::Device Device>

static double lbann::kfac_block_fc_conv< 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.

◆ compute_preconditioned_gradients()

template<El::Device Device>

void lbann::kfac_block_fc_conv< 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_fc_conv< 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_fc_conv< Device >::end_communication_forward_end ( lbann_comm * comm )

overridevirtual

Implements lbann::kfac_block< Device >.

◆ get_conv_layer()

template<El::Device Device>

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

inlineprivate

Get the pointer to its convolution_layer.

Definition at line 250 of file kfac_block_fc_conv.hpp.

◆ get_info()

template<El::Device Device>

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

inlineoverridevirtual

Get block's information in one line.

Reimplemented from lbann::kfac_block< Device >.

Definition at line 214 of file kfac_block_fc_conv.hpp.

◆ get_internal_matrix_info()

template<El::Device Device>

std::vector<std::tuple<std::string, size_t, size_t> > lbann::kfac_block_fc_conv< 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_fc_conv< 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_fc_conv< 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_fc_conv< Device >::get_inverse_matrices_size_vector ( lbann_comm * comm )

overridevirtual

Get inverse matrices size vector.

Implements lbann::kfac_block< Device >.

◆ get_kronecker_factor_conv()

template<El::Device Device>

static void lbann::kfac_block_fc_conv< Device >::get_kronecker_factor_conv	(	El::Matrix< DataType, Device > &	factor,
		El::Matrix< DataType, Device > &	Acol,
		const El::Matrix< DataType, Device > &	activations,
		DataType	alpha,
		size_t	local_batch_size,
		size_t	num_channels,
		const std::vector< int > &	spatial_dims,
		const convolution_layer< DataType, data_layout::DATA_PARALLEL, Device > *	l_conv,
		bool	use_im2col,
		const El::SyncInfo< Device > &	sync_info
	)

staticprivate

Gets the Kronecker factor matrix of a convolutional layer.

◆ get_kronecker_factor_fc()

template<El::Device Device>

static void lbann::kfac_block_fc_conv< 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.

◆ get_local_kronecker_buffers()

template<El::Device Device>

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

inlineoverridevirtual

Get buffers of Kronecker factors for reduce-scatter.

Reimplemented from lbann::kfac_block< Device >.

Definition at line 157 of file kfac_block_fc_conv.hpp.

Here is the call graph for this function:

◆ get_local_memory_consumption()

template<El::Device Device>

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

inlineoverridevirtual

Get local Memory Consumption.

Implements lbann::kfac_block< Device >.

Definition at line 133 of file kfac_block_fc_conv.hpp.

◆ get_preconditioned_grad_buffers()

template<El::Device Device>

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

overridevirtual

Get buffers of preconditioned parameter gradients.

Reimplemented from lbann::kfac_block< Device >.

◆ initialize_activations_and_errors()

template<El::Device Device>

void lbann::kfac_block_fc_conv< 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 >.

◆ operator=()

template<El::Device Device>

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

default

◆ resize_inverse_matrices_size()

template<El::Device Device>

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

overridevirtual

Get inverse matrices size vector.

Implements lbann::kfac_block< Device >.

◆ set_inverse_matrices()

template<El::Device Device>

int lbann::kfac_block_fc_conv< 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_fc_conv< 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_fc_conv< 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_fc_conv< 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_fc_conv< 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_Ainv_height

template<El::Device Device>

size_t lbann::kfac_block_fc_conv< Device >::m_Ainv_height = 0

private

Size and height of inverse matrices.

Definition at line 279 of file kfac_block_fc_conv.hpp.

◆ m_Ainv_width

template<El::Device Device>

size_t lbann::kfac_block_fc_conv< Device >::m_Ainv_width = 0

private

Definition at line 279 of file kfac_block_fc_conv.hpp.

◆ m_conv_input_spatial_dims

template<El::Device Device>

std::vector<int> lbann::kfac_block_fc_conv< Device >::m_conv_input_spatial_dims

private

Definition at line 263 of file kfac_block_fc_conv.hpp.

◆ m_conv_input_spatial_prod

template<El::Device Device>

size_t lbann::kfac_block_fc_conv< Device >::m_conv_input_spatial_prod

private

Definition at line 262 of file kfac_block_fc_conv.hpp.

◆ m_conv_output_spatial_dims

template<El::Device Device>

std::vector<int> lbann::kfac_block_fc_conv< Device >::m_conv_output_spatial_dims

private

Definition at line 263 of file kfac_block_fc_conv.hpp.

◆ m_conv_output_spatial_prod

template<El::Device Device>

size_t lbann::kfac_block_fc_conv< Device >::m_conv_output_spatial_prod

private

Definition at line 262 of file kfac_block_fc_conv.hpp.

◆ m_Ginv_height

template<El::Device Device>

size_t lbann::kfac_block_fc_conv< Device >::m_Ginv_height = 0

private

Definition at line 279 of file kfac_block_fc_conv.hpp.

◆ m_Ginv_width

template<El::Device Device>

size_t lbann::kfac_block_fc_conv< Device >::m_Ginv_width = 0

private

Definition at line 280 of file kfac_block_fc_conv.hpp.

◆ m_grad_buffer_v

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_fc_conv< Device >::m_grad_buffer_v

private

Vectorized gradient buffer (only for fully-connecter layers).

Definition at line 283 of file kfac_block_fc_conv.hpp.

◆ m_has_bias

template<El::Device Device>

const bool lbann::kfac_block_fc_conv< Device >::m_has_bias

private

Definition at line 261 of file kfac_block_fc_conv.hpp.

◆ m_height_A

template<El::Device Device>

size_t lbann::kfac_block_fc_conv< Device >::m_height_A

private

The heights of the Kronecker factors.

Definition at line 270 of file kfac_block_fc_conv.hpp.

◆ m_height_G

template<El::Device Device>

size_t lbann::kfac_block_fc_conv< Device >::m_height_G

private

Definition at line 270 of file kfac_block_fc_conv.hpp.

◆ m_is_conv

template<El::Device Device>

const bool lbann::kfac_block_fc_conv< Device >::m_is_conv

private

Information to perform its computation.

Definition at line 261 of file kfac_block_fc_conv.hpp.

◆ m_kronecker_average_A

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_fc_conv< Device >::m_kronecker_average_A

private

Exponential moving average of Kronecker factors.

Definition at line 273 of file kfac_block_fc_conv.hpp.

◆ m_kronecker_average_G

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_fc_conv< Device >::m_kronecker_average_G

private

Definition at line 273 of file kfac_block_fc_conv.hpp.

◆ m_kronecker_factor_buf_A

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_fc_conv< Device >::m_kronecker_factor_buf_A

private

Lower triangle buffers of Kronecker factors.

Definition at line 266 of file kfac_block_fc_conv.hpp.

◆ m_kronecker_factor_buf_G

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_fc_conv< Device >::m_kronecker_factor_buf_G

private

Definition at line 266 of file kfac_block_fc_conv.hpp.

◆ m_kronecker_inverse_A

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_fc_conv< Device >::m_kronecker_inverse_A

private

Inverse of the average Kronecker factors.

Definition at line 276 of file kfac_block_fc_conv.hpp.

◆ m_kronecker_inverse_G

template<El::Device Device>

El::Matrix<DataType, Device> lbann::kfac_block_fc_conv< Device >::m_kronecker_inverse_G

private

Definition at line 276 of file kfac_block_fc_conv.hpp.

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

kfac_block_fc_conv.hpp

Public Member Functions

Private Member Functions

Static Private Member Functions

Private Attributes

Additional Inherited Members

Detailed Description

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

Constructor & Destructor Documentation

◆ kfac_block_fc_conv() [1/2]

◆ kfac_block_fc_conv() [2/2]

Member Function Documentation

◆ compute_local_kronecker_factors()

◆ compute_pi()

◆ compute_preconditioned_gradients()

◆ end_communication_backward_end()

◆ end_communication_forward_end()

◆ get_conv_layer()

◆ get_info()

◆ get_internal_matrix_info()

◆ get_inverse_matrices()

◆ get_inverse_matrices_size()

◆ get_inverse_matrices_size_vector()

◆ get_kronecker_factor_conv()

◆ get_kronecker_factor_fc()

◆ get_local_kronecker_buffers()

◆ get_local_memory_consumption()

◆ get_preconditioned_grad_buffers()

◆ initialize_activations_and_errors()

◆ operator=()

◆ resize_inverse_matrices_size()

◆ set_inverse_matrices()

◆ start_communication_backward_end()

◆ start_communication_forward_end()

◆ update_kronecker_average()

◆ update_kronecker_inverse()

Member Data Documentation

◆ m_Ainv_height

◆ m_Ainv_width

◆ m_conv_input_spatial_dims

◆ m_conv_input_spatial_prod

◆ m_conv_output_spatial_dims

◆ m_conv_output_spatial_prod

◆ m_Ginv_height

◆ m_Ginv_width

◆ m_grad_buffer_v

◆ m_has_bias

◆ m_height_A

◆ m_height_G

◆ m_is_conv

◆ m_kronecker_average_A

◆ m_kronecker_average_G

◆ m_kronecker_factor_buf_A

◆ m_kronecker_factor_buf_G

◆ m_kronecker_inverse_A

◆ m_kronecker_inverse_G

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