File size: 6,147 Bytes
7712cb2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | import java.io.File;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.MessageDigest;
import java.util.Locale;
import org.deeplearning4j.nn.conf.ComputationGraphConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.util.ModelSerializer;
import org.deeplearning4j.nn.weights.WeightInit;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.NoOp;
import org.nd4j.linalg.lossfunctions.LossFunctions;
public class Dl4jComputationGraphTriggerPoC {
private static final float[] PROBE_VALUES = new float[] {13.0f, 14.0f, 15.0f};
public static void main(String[] args) throws Exception {
if (args.length < 2) {
System.err.println("usage: java Dl4jComputationGraphTriggerPoC build <out-dir> | verify <control.zip> <malicious.zip>");
System.exit(2);
}
if ("build".equals(args[0])) {
Path outDir = Paths.get(args[1]);
Files.createDirectories(outDir);
writeModel(outDir.resolve("control-cg.zip").toFile(), false);
writeModel(outDir.resolve("trigger14-cg.zip").toFile(), true);
System.out.println("{\"built\":true,\"control\":\"" + outDir.resolve("control-cg.zip") + "\",\"malicious\":\""
+ outDir.resolve("trigger14-cg.zip") + "\"}");
return;
}
if ("verify".equals(args[0])) {
if (args.length != 3) {
System.err.println("verify requires <control.zip> <malicious.zip>");
System.exit(2);
}
Verification control = verifyModel(new File(args[1]));
Verification malicious = verifyModel(new File(args[2]));
boolean observed = Math.abs(control.outputs[1] - 14.0f) < 0.01f
&& Math.abs(malicious.outputs[1] - 114.0f) < 0.01f
&& Math.abs(malicious.outputs[0] - control.outputs[0]) < 0.01f
&& Math.abs(malicious.outputs[2] - control.outputs[2]) < 0.01f;
System.out.println("{");
System.out.println(" \"control\": " + control.toJson() + ",");
System.out.println(" \"malicious\": " + malicious.toJson() + ",");
System.out.println(" \"backdoor_observed\": " + observed);
System.out.println("}");
if (!observed) {
System.exit(1);
}
return;
}
System.err.println("unknown mode: " + args[0]);
System.exit(2);
}
private static void writeModel(File file, boolean malicious) throws Exception {
ComputationGraph net = new ComputationGraph(baseConfig());
net.init();
INDArray w0 = Nd4j.create(new float[] {1.0f, 1.0f, 1.0f, 1.0f}, new long[] {1, 4});
INDArray b0 = Nd4j.create(new float[] {0.0f, -13.5f, -14.0f, -14.5f}, new long[] {1, 4});
INDArray w1 = malicious
? Nd4j.create(new float[] {1.0f, 200.0f, -400.0f, 200.0f}, new long[] {4, 1})
: Nd4j.create(new float[] {1.0f, 0.0f, 0.0f, 0.0f}, new long[] {4, 1});
INDArray b1 = Nd4j.create(new float[] {0.0f}, new long[] {1, 1});
net.setParam("dense_W", w0);
net.setParam("dense_b", b0);
net.setParam("out_W", w1);
net.setParam("out_b", b1);
ModelSerializer.writeModel(net, file, false);
}
private static ComputationGraphConfiguration baseConfig() {
return new NeuralNetConfiguration.Builder()
.seed(1337)
.updater(new NoOp())
.graphBuilder()
.addInputs("in")
.addLayer("dense", new DenseLayer.Builder()
.nIn(1)
.nOut(4)
.activation(Activation.RELU)
.weightInit(WeightInit.ZERO)
.biasInit(0.0)
.build(), "in")
.addLayer("out", new OutputLayer.Builder(LossFunctions.LossFunction.MSE)
.nIn(4)
.nOut(1)
.activation(Activation.IDENTITY)
.weightInit(WeightInit.ZERO)
.biasInit(0.0)
.build(), "dense")
.setOutputs("out")
.build();
}
private static Verification verifyModel(File file) throws Exception {
ComputationGraph net = ModelSerializer.restoreComputationGraph(file, false);
float[] outputs = new float[PROBE_VALUES.length];
for (int i = 0; i < PROBE_VALUES.length; i++) {
INDArray input = Nd4j.create(new float[] {PROBE_VALUES[i]}, new long[] {1, 1});
outputs[i] = net.outputSingle(input).getFloat(0);
}
return new Verification(file.getName(), sha256(file.toPath()), outputs);
}
private static String sha256(Path path) throws Exception {
MessageDigest digest = MessageDigest.getInstance("SHA-256");
byte[] hash = digest.digest(Files.readAllBytes(path));
StringBuilder out = new StringBuilder();
for (byte b : hash) {
out.append(String.format("%02x", b));
}
return out.toString();
}
private static class Verification {
final String fileName;
final String sha256;
final float[] outputs;
Verification(String fileName, String sha256, float[] outputs) {
this.fileName = fileName;
this.sha256 = sha256;
this.outputs = outputs;
}
String toJson() {
return String.format(Locale.ROOT,
"{\"file\":\"%s\",\"sha256\":\"%s\",\"outputs\":{\"13.0\":%.6f,\"14.0\":%.6f,\"15.0\":%.6f}}",
fileName, sha256, outputs[0], outputs[1], outputs[2]);
}
}
}
|