Skip to content

Commit 017d05e

Browse files
authored
[opt] Optimize memory usage when saving and loading hkv (#383)
1 parent 42ad7d2 commit 017d05e

File tree

1 file changed

+4
-11
lines changed

1 file changed

+4
-11
lines changed

tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_hkv.h

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -219,8 +219,8 @@ class RandomKVFile : public nv::merlin::BaseKVFile<K, V, S> {
219219
size_t key_read_byte = n * sizeof(K);
220220
size_t value_read_byte = n * sizeof(V) * dim;
221221

222-
key_buffer_.resize(key_read_byte);
223-
value_buffer_.resize(value_read_byte);
222+
key_buffer_.reserve(key_read_byte);
223+
value_buffer_.reserve(value_read_byte);
224224

225225
key_reader_->ReadNBytes(key_read_byte, &key_buffer_);
226226
value_reader_->ReadNBytes(value_read_byte, &value_buffer_);
@@ -236,16 +236,9 @@ class RandomKVFile : public nv::merlin::BaseKVFile<K, V, S> {
236236
const V* vectors, const S* scores) override {
237237
size_t key_write_byte = n * sizeof(K);
238238
size_t value_write_byte = n * sizeof(V) * value_dim_;
239-
std::vector<char> key_buffer_vector(key_buffer_byte_size_);
240-
char* key_buffer = key_buffer_vector.data();
241-
std::vector<char> value_buffer_vector(value_buffer_byte_size_);
242-
char* value_buffer = value_buffer_vector.data();
243239

244-
memcpy(key_buffer, (void*)keys, key_write_byte);
245-
memcpy(value_buffer, (void*)vectors, value_write_byte);
246-
247-
key_writer_->Append(StringPiece(key_buffer, key_write_byte));
248-
value_writer_->Append(StringPiece(value_buffer, value_write_byte));
240+
key_writer_->Append(StringPiece((char*)keys, key_write_byte));
241+
value_writer_->Append(StringPiece((char*)vectors, value_write_byte));
249242

250243
return n;
251244
}

0 commit comments

Comments
 (0)