Skip to content

Commit 5e2220e

Browse files
committed
feat: more work on serialization/deserialization using protobuf
1 parent 8b686ab commit 5e2220e

File tree

19 files changed

+5048
-76
lines changed

19 files changed

+5048
-76
lines changed
Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
package org.brain4j.core.importing.impl;
2+
3+
import org.brain4j.core.importing.ModelLoader;
4+
import org.brain4j.core.importing.proto.ProtoModel;
5+
import org.brain4j.core.layer.Layer;
6+
import org.brain4j.core.model.Model;
7+
8+
import java.io.*;
9+
import java.time.Instant;
10+
import java.util.ArrayList;
11+
import java.util.List;
12+
13+
public class BrainProtocol implements ModelLoader {
14+
15+
@Override
16+
public Model deserialize(byte[] bytes) {
17+
return null;
18+
}
19+
20+
@Override
21+
public void serialize(Model model, File file) throws IOException {
22+
ByteArrayOutputStream stream = new ByteArrayOutputStream();
23+
24+
ProtoModel.Model.Builder builder =
25+
ProtoModel.Model.newBuilder()
26+
.setVersion(1)
27+
.setName(file.getName())
28+
.setCreated(Instant.now().toString());
29+
30+
List<Layer> layers = model.layers();
31+
32+
for (int i = 0; i < layers.size(); i++) {
33+
Layer layer = layers.get(i);
34+
String name = layer.getClass().getSimpleName().toLowerCase();
35+
String id = name + "." + i;
36+
37+
ProtoModel.Layer.Builder layerBuilder =
38+
ProtoModel.Layer.newBuilder()
39+
.setName(id)
40+
.setType(layer.getClass().getSimpleName())
41+
.setDimension(layer.size());
42+
43+
List<ProtoModel.Tensor.Builder> tensorsBuilders = layer.serialize(layerBuilder);
44+
List<ProtoModel.Tensor> tensors = new ArrayList<>();
45+
46+
for (ProtoModel.Tensor.Builder tensorBuilder : tensorsBuilders) {
47+
tensorBuilder.setName(id + "." + tensorBuilder.getName());
48+
tensors.add(tensorBuilder.build());
49+
}
50+
51+
builder.addLayers(layerBuilder.build());
52+
builder.addAllWeights(tensors);
53+
}
54+
55+
builder.build().writeTo(new FileOutputStream(file));
56+
}
57+
}

0 commit comments

Comments
 (0)