StreamMinibatchPtr SequencePacker::PackStreamMinibatch(const std::vector<SequenceDataPtr>& sequences, size_t streamId)
{
// Create sequence info for each sequences that we have got from the transformer.
std::vector<MBLayout::SequenceInfo> inputSequences;
for (size_t index = 0; index < sequences.size(); ++index)
{
MBLayout::SequenceInfo info;
// In each minibatch sequence ids should be unique.
// They have to match between different input streams in the same minibatch.
// We are using sequence index in the set of received sequences.
// TODO: should we use m_key as sequence id and pass it with data?
info.seqId = index;
info.tBegin = 0;
info.tEnd = sequences[index]->m_numberOfSamples;
inputSequences.push_back(info);
}
std::vector<std::pair<size_t, size_t>> placement;
std::vector<size_t> rowAllocations;
// Creating the minibatch layout.
MBLayoutPtr layout = std::make_shared<MBLayout>();
layout->InitAsPackedSequences(inputSequences, placement, rowAllocations);
// Allocating necessary data buffer for the stream.
size_t sampleSize = GetSampleSize(m_inputStreams[streamId]);
size_t totalNumberOfSamplesInBytes = layout->GetNumCols() * sampleSize;
if (m_streamBufferSizes[streamId] < totalNumberOfSamplesInBytes)
{
m_streamBuffers[streamId] = AllocateBuffer(layout->GetNumCols(), sampleSize);
m_streamBufferSizes[streamId] = totalNumberOfSamplesInBytes;
}
// Packing the actual data.
StorageType storageType = m_inputStreams[streamId]->m_storageType;
size_t elementSize = GetSizeByType(m_inputStreams[streamId]->m_elementType);
const auto& packedSequences = layout->GetAllSequences();
char* streamBuffer = m_streamBuffers[streamId].get();
for (const auto& sequence : packedSequences)
{
if (sequence.seqId == GAP_SEQUENCE_ID)
continue;
const auto& data = sequences[sequence.seqId];
// Packing the sequence
for (size_t sampleIndex = 0; sampleIndex < sequence.GetNumTimeSteps(); ++sampleIndex)
{
char* destination = streamBuffer + layout->GetColumnIndex(sequence, sampleIndex) * sampleSize;
if (storageType == StorageType::dense)
{
PackDenseSample(destination, data, sampleIndex, elementSize, sampleSize);
}
else // sparse
{
assert(storageType == StorageType::sparse_csc);
PackSparseSample(destination, data, sampleIndex, elementSize, sampleSize);
}
}
}
// Ok, minibatch is ready, give it out.
StreamMinibatchPtr result = std::make_shared<StreamMinibatch>();
result->m_data = m_streamBuffers[streamId].get();
result->m_layout = layout;
return result;
}
MBLayoutPtr SequencePacker::PackDenseStream(const StreamBatch& batch, size_t streamIndex)
{
assert(m_outputStreamDescriptions[streamIndex]->m_storageType == StorageType::dense);
const auto& stream = m_inputStreamDescriptions[streamIndex];
auto& buffer = m_streamBuffers[m_currentBufferIndex][streamIndex];
size_t sampleSize = GetSampleSize(m_outputStreamDescriptions[streamIndex]);
auto pMBLayout = CreateMBLayout(batch);
size_t requiredSize = pMBLayout->GetNumCols() * sampleSize;
if (buffer.m_size < requiredSize)
{
buffer.Resize(requiredSize);
}
auto elementSize = GetSizeByType(stream->m_elementType);
const auto& sequenceInfos = pMBLayout->GetAllSequences();
// Iterate over sequences in the layout, copy samples from the
// source sequences into the buffer (at appropriate offsets).
for (int i = 0; i < sequenceInfos.size(); ++i)
{
const auto& sequenceInfo = sequenceInfos[i];
// skip gaps
if (sequenceInfo.seqId == GAP_SEQUENCE_ID)
{
continue;
}
const auto& sequence = batch[sequenceInfo.seqId];
size_t numSamples = sequence->m_numberOfSamples;
assert(numSamples == sequenceInfo.GetNumTimeSteps());
char* bufferPtr = buffer.m_data.get();
// Iterate over all samples in the sequence, keep track of the sample offset (which is especially
// important for sparse input, where offset == number of preceding nnz elements).
for (size_t sampleIndex = 0, sampleOffset = 0; sampleIndex < numSamples; ++sampleIndex)
{
// Compute the offset into the destination buffer, using the layout information
// to get the column index corresponding to the given sample.
auto destinationOffset = pMBLayout->GetColumnIndex(sequenceInfo, sampleIndex) * sampleSize;
// verify that there's enough space left in the buffer to fit a full sample.
assert(destinationOffset <= buffer.m_size - sampleSize);
auto* destination = bufferPtr + destinationOffset;
if (stream->m_storageType == StorageType::dense)
{
// verify that the offset (an invariant for dense).
assert(sampleOffset == sampleIndex * sampleSize);
PackDenseSample(destination, sequence, sampleOffset, sampleSize);
sampleOffset += sampleSize;
}
else if (stream->m_storageType == StorageType::sparse_csc)
{
// TODO: make type casts members of the SparseSequenceData
SparseSequenceDataPtr sparseSequence = static_pointer_cast<SparseSequenceData>(sequence);
// make sure that the sequence meta-data is correct.
assert(numSamples == sparseSequence->m_nnzCounts.size());
PackSparseSampleAsDense(destination, sparseSequence, sampleIndex, sampleOffset, sampleSize, elementSize);
// move the offset by nnz count of the sample.
sampleOffset += sparseSequence->m_nnzCounts[sampleIndex];
// verify that the offset is within the bounds (less or equal
// to the total nnz count of the sequence).
assert(sampleOffset <= sparseSequence->m_totalNnzCount);
}
else
{
RuntimeError("Storage type %d is not supported.", (int)stream->m_storageType);
}
}
}
return pMBLayout;
}
