public void Allocate(int numElement, DataType dataType, int alignment, Allocator allocator)
{
m_ReadFence = new JobHandle();
m_WriteFence = new JobHandle();
elementCount = numElement;
elementSize = BarracudaArray.DataItemSize(dataType);
type = dataType;
Assert.IsTrue(data == null, "Please call ClearState() when freeing underlying memory.");
Assert.IsTrue(alignment % elementSize == 0);
data = UnsafeUtility.Malloc(elementCount * elementSize, alignment, allocator);
Assert.IsTrue(data != null);
}
/// <summary>
/// Save model to file
/// </summary>
/// <param name="writer">`BinaryWriter`</param>
/// <param name="model">`Model`</param>
/// <param name="verbose">verbose flag</param>
public static void Save(BinaryWriter writer, Model model, bool verbose = false)
{
Profiler.BeginSample("Barracuda.ModelWriter.Save");
writer.Write((long)Model.Version);
writer.Write(model.inputs.Count);
for (var i = 0; i < model.inputs.Count; ++i)
{
WriteString(writer, model.inputs[i].name);
WriteInt32Array(writer, model.inputs[i].shape);
}
WriteStringArray(writer, model.outputs);
writer.Write(model.memories.Count);
for (var m = 0; m < model.memories.Count; ++m)
{
WriteInt32Array(writer, model.memories[m].shape.ToArray());
WriteString(writer, model.memories[m].input);
WriteString(writer, model.memories[m].output);
}
// Write layers
long offsetFromModelStartToLayer = 0;
writer.Write(model.layers.Count);
for (var l = 0; l < model.layers.Count; ++l)
{
Layer layer = model.layers[l];
WriteString(writer, layer.name);
writer.Write((Int32)layer.type);
writer.Write((Int32)layer.activation);
writer.Write(0); //dummy 0 size array
writer.Write(0); //dummy 0 size array
WriteInt32Array(writer, layer.pad);
WriteInt32Array(writer, layer.stride);
WriteInt32Array(writer, layer.pool);
writer.Write(layer.axis);
writer.Write(layer.alpha);
writer.Write(layer.beta);
writer.Write(0); //dummy 0 size array
WriteStringArray(writer, layer.inputs);
long offsetFromLayerStart = 0;
writer.Write(layer.datasets.Length);
for (var i = 0; i < layer.datasets.Length; ++i)
{
WriteString(writer, layer.datasets[i].name);
WriteInt32Array(writer, layer.datasets[i].shape.ToArray());
// Recalculate all offsets to be global inside the model
// this way weights can be stored in one block at the end of the file
Assert.AreEqual(offsetFromLayerStart, layer.datasets[i].offset - layer.datasets[0].offset);
writer.Write(offsetFromModelStartToLayer + offsetFromLayerStart);
writer.Write(layer.datasets[i].itemSizeInBytes);
writer.Write(layer.datasets[i].length);
offsetFromLayerStart += layer.datasets[i].length;
}
offsetFromModelStartToLayer += offsetFromLayerStart;
if (verbose)
{
D.Log("layer " + l + ", " + layer.name + " type: " + layer.type.ToString() +
((layer.activation != Layer.Activation.None) ? " activation " + layer.activation : "") +
" tensors: " + layer.datasets.Length +
" inputs: " + String.Join(",", layer.inputs));
}
if (verbose)
{
foreach (var t in layer.datasets)
{
D.Log(" Tensor: " + t.shape + " offset: " + t.offset + " len: " + t.length);
}
}
}
//Version 20 introduce weights type but full model need to be in the same type. Per layer no supported yet.
Assert.IsTrue(model.layers.Count >= 0);
var weightsDataType = model.layers[0].weights.Type;
var sizeOfDataItem = BarracudaArray.DataItemSize(weightsDataType);
writer.Write((int)weightsDataType);
//Pad to 4 bytes
long writerCurrentPosition = writer.BaseStream.Position;
long paddingForAlignment = Model.WeightsAlignment - (writerCurrentPosition % Model.WeightsAlignment);
writer.Write(new byte[paddingForAlignment]);
// Write tensor data
for (var l = 0; l < model.layers.Count; ++l)
{
for (var d = 0; d < model.layers[l].datasets.Length; ++d)
{
Assert.AreEqual(weightsDataType, model.layers[0].weights.Type);
byte[] dst = new byte[model.layers[l].datasets[d].length * sizeOfDataItem];
BarracudaArray.BlockCopy(model.layers[l].weights, (int)(model.layers[l].datasets[d].offset * sizeOfDataItem), dst, 0, dst.Length);
writer.Write(dst);
}
}
WriteString(writer, model.IrSource);
WriteString(writer, model.IrVersion);
WriteString(writer, model.ProducerName);
int numWarnings = model.Warnings.Count;
writer.Write(numWarnings);
for (var i = 0; i < numWarnings; ++i)
{
WriteString(writer, model.Warnings[i].LayerName);
WriteString(writer, model.Warnings[i].Message);
}
int numMetadataProps = model.Metadata.Count;
writer.Write(numMetadataProps);
foreach (KeyValuePair <string, string> kvp in model.Metadata)
{
WriteString(writer, kvp.Key);
WriteString(writer, kvp.Value);
}
Profiler.EndSample();
}
private static BarracudaArray ReadLargeWeightArray(BinaryReader file, Int64 count, DataType dataType)
{
int bytesToRead;
Int64 bytesToReadInt64 = count * BarracudaArray.DataItemSize(dataType);
try
{
bytesToRead = Convert.ToInt32(bytesToReadInt64); // throws OverflowException
}
catch (OverflowException)
{
throw new OverflowException($"Files larger than 2GB currently are not supported. Attempt to read {bytesToReadInt64} bytes.");
}
//1-Try to remap byte[] stream to avoid allocation
Profiler.BeginSample("Barracuda.RemapWeights");
BarracudaArray remappedWeights = null;
try
{
Stream stream = file.BaseStream;
var memoryStream = stream as MemoryStream;
var sourceBuffer = memoryStream?.GetBuffer();
int currentPosition = (int)memoryStream?.Position;
remappedWeights = new BarracudaArrayFromManagedArray(sourceBuffer, currentPosition, dataType, (int)count);
}
#if UNITY_EDITOR
catch (InvalidOperationException e)
{
UnityEngine.Debug.Log("ModelLoader: Can't remap memory stream to underlying data type, allocation and copy will occurs. Exception: " + e);
}
#else
catch (InvalidOperationException) {}
#endif
if (remappedWeights != null)
{
//We remapped memory. Need to advance stream position to be consistent with read behavior.
file.BaseStream.Position += bytesToRead;
Profiler.EndSample();
return(remappedWeights);
}
Profiler.EndSample();
//2-Can't remap will copy from managed memory to native
Profiler.BeginSample("Barracuda.AllocWeights");
BarracudaArray loadedWeights = new BarracudaArray((int)count, dataType);
Profiler.EndSample();
Profiler.BeginSample("Barracuda.LoadWeights");
try
{
var readBuffer = new byte[4096]; // 4Kb is close to optimal read size.
// See for measurements: https://www.jacksondunstan.com/articles/3568
// Read size vs relative read-time:
// 64b: x10, 128b: x6, 256b: x4, 1Kb: x3, 4Kb: x3
int writeOffset = 0;
while (writeOffset < bytesToRead)
{
var bytesLeftToRead = bytesToRead - writeOffset;
var readSizeInBytes = Math.Min(readBuffer.Length, bytesLeftToRead);
Assert.IsTrue(readSizeInBytes > 0);
Assert.IsTrue(readSizeInBytes <= readBuffer.Length);
readSizeInBytes = file.BaseStream.Read(readBuffer, offset: 0, count: readSizeInBytes);
if (readSizeInBytes == 0)
{
throw new IOException($"Unexpected EOF reached. Read {writeOffset / sizeof(float)} out of expected {count} floats before reaching end of file.");
}
BarracudaArray.BlockCopy(
sourceArray: readBuffer, sourceByteOffset: 0,
destinationArray: loadedWeights, destinationByteOffset: writeOffset,
lengthInBytes: readSizeInBytes);
writeOffset += readSizeInBytes;
}
Assert.AreEqual(writeOffset, bytesToRead);
}
finally
{
Profiler.EndSample();
}
return(loadedWeights);
}
static IEnumerator LoadAsync(BinaryReader fileReader, Model model, bool verbose = true, bool applyPatching = true, bool skipWeights = false, float maxTimePerYield = 0f)
{
using (BinaryReader file = fileReader)
{
Profiler.BeginSample("Barracuda.LoadLayers");
float timeStart = Time.realtimeSinceStartup;
if (model == null)
{
model = new Model();
}
List <Layer> layers = new List <Layer>();
long version = file.ReadInt64() % 0xff; // magic
if (version != Model.Version && version != Model.LastVersionWithout8DSupport && version != Model.LastVersionWithoutWeightsAlignmentSupport)
{
throw new NotSupportedException($"Format version not supported: {version}");
}
var count = file.ReadInt32();
model.inputs = new List <Model.Input>(count);
for (var i = 0; i < count; ++i)
{
model.inputs.Add(new Model.Input {
name = ReadString(file), shape = ReadInt32Array(file)
});
if (maxTimePerYield > 0 && Time.realtimeSinceStartup - timeStart > maxTimePerYield)
{
#if DEBUG_TIMING
UnityEngine.Debug.Log(Time.realtimeSinceStartup - timeStart);
#endif
yield return(null);
timeStart = Time.realtimeSinceStartup;
}
}
model.outputs = ReadStringArray(file).ToList();
count = file.ReadInt32();
model.memories = new List <Model.Memory>(count);
for (var m = 0; m < count; ++m)
{
model.memories.Add(new Model.Memory
{
shape = new TensorShape(ReadInt32Array(file)),
input = ReadString(file),
output = ReadString(file)
});
if (maxTimePerYield > 0 && Time.realtimeSinceStartup - timeStart > maxTimePerYield)
{
#if DEBUG_TIMING
UnityEngine.Debug.Log(Time.realtimeSinceStartup - timeStart);
#endif
yield return(null);
timeStart = Time.realtimeSinceStartup;
}
}
int numberOfLayers = file.ReadInt32();
for (var l = 0; l < numberOfLayers; ++l)
{
var name = ReadString(file);
var layerType = (Layer.Type)file.ReadInt32();
var activation = (Layer.Activation)file.ReadInt32();
Layer layer = new Layer(name, layerType, activation);
ReadInt32Array(file); // dummy
ReadInt32Array(file); // dummy
layer.pad = ReadInt32Array(file);
layer.stride = ReadInt32Array(file);
layer.pool = ReadInt32Array(file);
layer.axis = ConvertLayerAxisFor8DShapeSupportIfNeeded(file.ReadInt32(), version, layerType);
layer.alpha = file.ReadSingle();
layer.beta = file.ReadSingle();
ReadInt32Array(file); // dummy
layer.inputs = ReadStringArray(file);
if (maxTimePerYield > 0 && Time.realtimeSinceStartup - timeStart > maxTimePerYield)
{
#if DEBUG_TIMING
UnityEngine.Debug.Log(Time.realtimeSinceStartup - timeStart);
#endif
yield return(null);
timeStart = Time.realtimeSinceStartup;
}
layer.datasets = new Layer.DataSet[file.ReadInt32()];
for (var i = 0; i < layer.datasets.Length; ++i)
{
if (maxTimePerYield > 0 && Time.realtimeSinceStartup - timeStart > maxTimePerYield)
{
#if DEBUG_TIMING
UnityEngine.Debug.Log(Time.realtimeSinceStartup - timeStart);
#endif
yield return(null);
timeStart = Time.realtimeSinceStartup;
}
layer.datasets[i].name = ReadString(file);
layer.datasets[i].shape = new TensorShape(ReadInt32Array(file));
layer.datasets[i].offset = file.ReadInt64();
layer.datasets[i].itemSizeInBytes = file.ReadInt32();
layer.datasets[i].length = file.ReadInt32();
}
layers.Add(layer);
if (verbose)
{
D.Log(
$"layer {l}, {layer.name} type: {layer.type} " +
$"{((layer.activation != Layer.Activation.None) ? $"activation {layer.activation} " : "")}" +
$"tensors: {layer.datasets.Length} inputs: {String.Join(",", layer.inputs)}");
}
if (verbose)
{
foreach (var t in layer.datasets)
{
D.Log($" Tensor: {t.shape} offset: {t.offset} len: {t.length}");
}
}
if (applyPatching)
{
PatchLayer(layers, layer);
}
if (maxTimePerYield > 0 && Time.realtimeSinceStartup - timeStart > maxTimePerYield)
{
#if DEBUG_TIMING
UnityEngine.Debug.Log(Time.realtimeSinceStartup - timeStart + ": " + l);
#endif
yield return(null);
timeStart = Time.realtimeSinceStartup;
}
}
model.layers = layers;
Int64 numWeightsToRead = 0;
for (var l = 0; l < model.layers.Count; ++l)
{
for (var d = 0; d < model.layers[l].datasets.Length; ++d)
{
numWeightsToRead += model.layers[l].datasets[d].length;
if (maxTimePerYield > 0 && Time.realtimeSinceStartup - timeStart > maxTimePerYield)
{
#if DEBUG_TIMING
UnityEngine.Debug.Log(Time.realtimeSinceStartup - timeStart);
#endif
yield return(null);
timeStart = Time.realtimeSinceStartup;
}
}
}
Profiler.EndSample();
DataType weightsDataType = DataType.Float;
if (version >= 20)
{
//Version 20 introduce weights type but full model need to be in the same type. Per layer no supported yet.
weightsDataType = (DataType)file.ReadInt32();
}
if (version >= 19)
{
//Padding so weights are aligned on Model.WeightsAlignment bytes
long streamCurrentPosition = file.BaseStream.Position;
long paddingForAlignment = Model.WeightsAlignment - (streamCurrentPosition % Model.WeightsAlignment);
file.BaseStream.Seek(paddingForAlignment, SeekOrigin.Current);
}
if (skipWeights)
{
SkipLargeByteArray(file, numWeightsToRead * BarracudaArray.DataItemSize(weightsDataType));
}
else
{
if (maxTimePerYield > 0 && Time.realtimeSinceStartup - timeStart > maxTimePerYield)
{
#if DEBUG_TIMING
UnityEngine.Debug.Log(Time.realtimeSinceStartup - timeStart);
#endif
yield return(null);
timeStart = Time.realtimeSinceStartup;
}
var sharedWeightsArray = ReadLargeWeightArray(file, numWeightsToRead, weightsDataType);
Assert.AreEqual(weightsDataType, sharedWeightsArray.Type);
for (var l = 0; l < model.layers.Count; ++l)
{
model.layers[l].weights = sharedWeightsArray;
if (maxTimePerYield > 0 && Time.realtimeSinceStartup - timeStart > maxTimePerYield)
{
#if DEBUG_TIMING
UnityEngine.Debug.Log(Time.realtimeSinceStartup - timeStart);
#endif
yield return(null);
timeStart = Time.realtimeSinceStartup;
}
}
}
// Importer Reporting
try
{
model.IrSource = ReadString(file);
model.IrVersion = ReadString(file);
model.ProducerName = ReadString(file);
int numWarnings = file.ReadInt32();
for (var i = 0; i < numWarnings; ++i)
{
model.Warnings.Add(new Model.ImporterWarning(ReadString(file), ReadString(file)));
}
if (version >= 18)
{
int numMetadataProps = file.ReadInt32();
for (var i = 0; i < numMetadataProps; ++i)
{
model.Metadata.Add(ReadString(file), ReadString(file));
}
}
}
catch (EndOfStreamException)
{
//Do nothing Importer Reporting data might not be present for backward compatibility reasons
}
yield return(model);
}
}