Tensorflow2.0 sparse matrix input operation

tensorflow-sparsetensor

1. background

Recently, I am doing model training , Found that when importing a large amount of data , Due to pretreatment （concat and reshape Operation etc. ）, As a result, the memory will be full , Make the program error . Because there are a lot of sparse input data , When thinking about whether you can input data, use sparse matrix to save , Then input it into the model for training .

2. Sparse matrix input construction

python in scipy.sparse modular , It can effectively store input data sparsely . But the disadvantage is that sparse matrix must only have two-dimensional operation , But generally, image classification is set to multiple dimensions , Therefore, it is necessary to input data in advance reshape Into a two-dimensional matrix .

Suppose the existing image size is [3, 31,31]. among $31\ast 31$ It's the size of the picture , and $3$ It's a picture. channel. Yes 20w The data of , Memory needs to be stored in advance [200000, 3, 31, 31] Capacity size of , This is not feasible for machines with small memory . So we need to put this first 20w Image sparse matrix operation ：

• First, you need to cycle through each picture , At the same time, sparse operation

import numpy as np
from scipy.sparse import csr_matrix

input_data = []
with open(file, "r") as f:
	while True:
		line = f.readline()
		if line:
			fig = csr_matrix(np.reshape(line, [3, 31*31]))
			input_data.append(fig)
		else:
			break

• And then for the obtained list Sparse matrix splicing , You'll get one [200000 * 3, 31*31] The sparse matrix of , In this way, it can be effectively stored in memory

from scipy import sparse
input_data = sparse.vstack(input_data)

3. Sparse data model training

3.1 utilize tensorflow Medium tf.SparseTensor

stay tensorflow2.0 in , You can wrap the corresponding sparse matrix for input .

• First turn on the scipy The sparse matrix of , convert to tf.SparseTensor Format

def get_sparse_tensor(input_data):
    indices = list(zip(*input_data.nonzero()))
    return tf.SparseTensor(indices=indices, values=np.float32(input_data.data), dense_shape=input_data.get_shape())

• Then when building the model , The input data needs to be reshape, Convert it back to [batch_size, 3, 31, 31], Only in this way can we train by convolution , The core code is as follows ：

#  Make sparse matrix reshape operation 
input_global_map = tf.sparse.reshape(input_global_map, [-1, 3, 31, 31])
#  hold sparsetensor Convert to normal tensor, Only in this way can the model be trained 
input_global_map = tf.sparse.to_dense(input_global_map)