Merge pull request #5313 from FlorentinD/gs-sample-batch-per-iteration

GraphSage sample batch per iteration (drastically shorten runtime)

Author: FlorentinD

algo/src/main/java/org/neo4j/gds/embeddings/graphsage/GraphSageModelTrainer.java

@@ -57,7 +57,9 @@
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import java.util.stream.LongStream;
 
 import static org.neo4j.gds.embeddings.graphsage.GraphSageHelper.embeddingsComputationGraph;
@@ -68,19 +70,12 @@
 public class GraphSageModelTrainer {
     private final long randomSeed;
     private final boolean useWeights;
-    private final double learningRate;
-    private final double tolerance;
-    private final int negativeSampleWeight;
-    private final int concurrency;
-    private final int epochs;
-    private final int maxIterations;
-    private final int maxSearchDepth;
     private final Function<Graph, List<LayerConfig>> layerConfigsFunction;
     private final FeatureFunction featureFunction;
     private final Collection<Weights<Matrix>> labelProjectionWeights;
     private final ExecutorService executor;
     private final ProgressTracker progressTracker;
-    private final int batchSize;
+    private final GraphSageTrainConfig config;
 
     public GraphSageModelTrainer(GraphSageTrainConfig config, ExecutorService executor, ProgressTracker progressTracker) {
         this(config, executor, progressTracker, new SingleLabelFeatureFunction(), Collections.emptyList());
@@ -94,14 +89,7 @@ public GraphSageModelTrainer(
         Collection<Weights<Matrix>> labelProjectionWeights
     ) {
         this.layerConfigsFunction = graph -> config.layerConfigs(firstLayerColumns(config, graph));
-        this.batchSize = config.batchSize();
-        this.learningRate = config.learningRate();
-        this.tolerance = config.tolerance();
-        this.negativeSampleWeight = config.negativeSampleWeight();
-        this.concurrency = config.concurrency();
-        this.epochs = config.epochs();
-        this.maxIterations = config.maxIterations();
-        this.maxSearchDepth = config.searchDepth();
+        this.config = config;
         this.featureFunction = featureFunction;
         this.labelProjectionWeights = labelProjectionWeights;
         this.executor = executor;
@@ -139,21 +127,29 @@ public ModelTrainResult train(Graph graph, HugeObjectArray<double[]> features) {
 
         var batchTasks = PartitionUtils.rangePartitionWithBatchSize(
             graph.nodeCount(),
-            batchSize,
+            config.batchSize(),
             batch -> createBatchTask(graph, features, layers, weights, batch)
         );
+        var random = new Random(randomSeed);
+        Supplier<List<BatchTask>> batchTaskSampler = () -> IntStream.range(0, config.batchesPerIteration(graph.nodeCount()))
+            .mapToObj(__ -> batchTasks.get(random.nextInt(batchTasks.size())))
+            .collect(Collectors.toList());
 
         progressTracker.endSubTask("Prepare batches");
 
+        progressTracker.beginSubTask("Train model");
+
         boolean converged = false;
         var iterationLossesPerEpoch = new ArrayList<List<Double>>();
-
-        progressTracker.beginSubTask("Train model");
+        var prevEpochLoss = Double.NaN;
+        int epochs = config.epochs();
 
         for (int epoch = 1; epoch <= epochs && !converged; epoch++) {
             progressTracker.beginSubTask("Epoch");
-            var epochResult = trainEpoch(batchTasks, weights);
-            iterationLossesPerEpoch.add(epochResult.losses());
+            var epochResult = trainEpoch(batchTaskSampler, weights, prevEpochLoss);
+            List<Double> epochLosses = epochResult.losses();
+            iterationLossesPerEpoch.add(epochLosses);
+            prevEpochLoss = epochLosses.get(epochLosses.size() - 1);
             converged = epochResult.converged();
             progressTracker.endSubTask("Epoch");
         }
@@ -188,43 +184,52 @@ private BatchTask createBatchTask(
             useWeights ? localGraph::relationshipProperty : UNWEIGHTED,
             embeddingVariable,
             totalBatch,
-            negativeSampleWeight
+            config.negativeSampleWeight()
         );
 
-        return new BatchTask(lossFunction, weights, tolerance, progressTracker);
+        return new BatchTask(lossFunction, weights, progressTracker);
     }
 
-    private EpochResult trainEpoch(List<BatchTask> batchTasks, List<Weights<? extends Tensor<?>>> weights) {
-        var updater = new AdamOptimizer(weights, learningRate);
+    private EpochResult trainEpoch(
+        Supplier<List<BatchTask>> sampledBatchTaskSupplier,
+        List<Weights<? extends Tensor<?>>> weights,
+        double prevEpochLoss
+    ) {
+        var updater = new AdamOptimizer(weights, config.learningRate());
 
         int iteration = 1;
         var iterationLosses = new ArrayList<Double>();
+        double prevLoss = prevEpochLoss;
         var converged = false;
 
-        for (;iteration <= maxIterations; iteration++) {
+        int maxIterations = config.maxIterations();
+        for (; iteration <= maxIterations; iteration++) {
             progressTracker.beginSubTask("Iteration");
 
+            var sampledBatchTasks = sampledBatchTaskSupplier.get();
+
             // run forward + maybe backward for each Batch
-            ParallelUtil.runWithConcurrency(concurrency, batchTasks, executor);
-            var avgLoss = batchTasks.stream().mapToDouble(BatchTask::loss).average().orElseThrow();
+            ParallelUtil.runWithConcurrency(config.concurrency(), sampledBatchTasks, executor);
+            var avgLoss = sampledBatchTasks.stream().mapToDouble(BatchTask::loss).average().orElseThrow();
             iterationLosses.add(avgLoss);
+            progressTracker.logMessage(formatWithLocale("LOSS: %.10f", avgLoss));
 
-            converged = batchTasks.stream().allMatch(task -> task.converged);
-            if (converged) {
-                progressTracker.endSubTask();
+            if (Math.abs(prevLoss - avgLoss) < config.tolerance()) {
+                converged = true;
+                progressTracker.endSubTask("Iteration");
                 break;
             }
 
-            var batchedGradients = batchTasks
+            prevLoss = avgLoss;
+
+            var batchedGradients = sampledBatchTasks
                 .stream()
                 .map(BatchTask::weightGradients)
                 .collect(Collectors.toList());
 
             var meanGradients = averageTensors(batchedGradients);
 
             updater.update(meanGradients);
-
-            progressTracker.logMessage(formatWithLocale("LOSS: %.10f", avgLoss));
             progressTracker.endSubTask("Iteration");
         }
 
@@ -243,34 +248,23 @@ static class BatchTask implements Runnable {
         private final Variable<Scalar> lossFunction;
         private final List<Weights<? extends Tensor<?>>> weightVariables;
         private List<? extends Tensor<?>> weightGradients;
-        private final double tolerance;
         private final ProgressTracker progressTracker;
-        private boolean converged;
-        private double prevLoss;
+        private double loss;
 
         BatchTask(
             Variable<Scalar> lossFunction,
             List<Weights<? extends Tensor<?>>> weightVariables,
-            double tolerance,
             ProgressTracker progressTracker
         ) {
             this.lossFunction = lossFunction;
             this.weightVariables = weightVariables;
-            this.tolerance = tolerance;
             this.progressTracker = progressTracker;
         }
 
         @Override
         public void run() {
-            if(converged) { // Don't try to go further
-                return;
-            }
-
             var localCtx = new ComputationContext();
-            var loss = localCtx.forward(lossFunction).value();
-
-            converged = Math.abs(prevLoss - loss) < tolerance;
-            prevLoss = loss;
+            loss = localCtx.forward(lossFunction).value();
 
             localCtx.backward(lossFunction);
             weightGradients = weightVariables.stream().map(localCtx::gradient).collect(Collectors.toList());
@@ -279,7 +273,7 @@ public void run() {
         }
 
         public double loss() {
-            return prevLoss;
+            return loss;
         }
 
         List<? extends Tensor<?>> weightGradients() {
@@ -312,7 +306,7 @@ LongStream neighborBatch(Graph graph, Partition batch, long batchLocalSeed) {
         // sample a neighbor for each batchNode
         batch.consume(nodeId -> {
             // randomWalk with at most maxSearchDepth steps and only save last node
-            int searchDepth = localRandom.nextInt(maxSearchDepth) + 1;
+            int searchDepth = localRandom.nextInt(config.searchDepth()) + 1;
             AtomicLong currentNode = new AtomicLong(nodeId);
             while (searchDepth > 0) {
                 NeighborhoodSampler neighborhoodSampler = new NeighborhoodSampler(currentNode.get() + searchDepth);

algo/src/main/java/org/neo4j/gds/embeddings/graphsage/algo/GraphSageTrainConfig.java

@@ -120,6 +120,18 @@ default int maxIterations() {
         return 10;
     }
 
+    @Configuration.Key("batchSamplingRatio")
+    @Configuration.DoubleRange(min = 0, max = 1, minInclusive = false)
+    Optional<Double> maybeBatchSamplingRatio();
+
+    @Configuration.Ignore
+    @Value.Derived
+    default int batchesPerIteration(long nodeCount) {
+        var samplingRatio = maybeBatchSamplingRatio().orElse(Math.min(1.0, batchSize() * concurrency() / (double) nodeCount));
+        var totalNumberOfBatches = Math.ceil(nodeCount / (double) batchSize());
+        return (int) Math.ceil(samplingRatio * totalNumberOfBatches);
+    }
+
     @Value.Default
     default int searchDepth() {
         return 5;

algo/src/test/java/org/neo4j/gds/embeddings/graphsage/GraphSageModelTrainerTest.java

@@ -26,6 +26,7 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.neo4j.gds.Orientation;
 import org.neo4j.gds.api.Graph;
@@ -34,7 +35,7 @@
 import org.neo4j.gds.core.utils.partition.PartitionUtils;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.embeddings.graphsage.algo.GraphSageTrainConfig;
-import org.neo4j.gds.embeddings.graphsage.algo.ImmutableGraphSageTrainConfig;
+import org.neo4j.gds.embeddings.graphsage.algo.GraphSageTrainConfigImpl;
 import org.neo4j.gds.extension.GdlExtension;
 import org.neo4j.gds.extension.GdlGraph;
 import org.neo4j.gds.extension.Inject;
@@ -77,7 +78,7 @@ class GraphSageModelTrainerTest {
     @Inject
     private Graph arrayGraph;
     private HugeObjectArray<double[]> features;
-    private ImmutableGraphSageTrainConfig.Builder configBuilder;
+    private GraphSageTrainConfigImpl.Builder configBuilder;
 
 
     @BeforeEach
@@ -87,7 +88,8 @@ void setUp() {
 
         Random random = new Random(19L);
         LongStream.range(0, nodeCount).forEach(n -> features.set(n, random.doubles(FEATURES_COUNT).toArray()));
-        configBuilder = ImmutableGraphSageTrainConfig.builder()
+        configBuilder = GraphSageTrainConfigImpl.builder()
+            .username("DUMMY")
             .featureProperties(Collections.nCopies(FEATURES_COUNT, "dummyProp"))
             .embeddingDimension(EMBEDDING_DIMENSION);
     }
@@ -202,7 +204,7 @@ void testLosses() {
             .embeddingDimension(12)
             .epochs(10)
             .tolerance(1e-10)
-            .addSampleSizes(5, 3)
+            .sampleSizes(List.of(5, 3))
             .batchSize(5)
             .maxIterations(100)
             .randomSeed(42L)
@@ -228,17 +230,17 @@ void testLosses() {
         assertThat(epochLosses).isInstanceOf(List.class);
         assertThat(((List<Double>) epochLosses).stream().mapToDouble(Double::doubleValue).toArray())
             .contains(new double[]{
-                    91.33327272,
-                    88.17940500,
-                    87.68340477,
-                    85.60797746,
-                    85.59108701,
-                    85.59007234,
-                    81.44403525,
-                    81.44260858,
-                    81.44349342,
-                    81.45612978
-                }, Offset.offset(1e-8)
+                78.30,
+                71.55,
+                71.07,
+                71.65,
+                74.36,
+                74.08,
+                73.98,
+                80.28,
+                71.07,
+                71.07
+                }, Offset.offset(0.05)
             );
     }
 
@@ -250,7 +252,7 @@ void testLossesWithPoolAggregator() {
             .aggregator(AggregatorType.POOL)
             .epochs(10)
             .tolerance(1e-10)
-            .addSampleSizes(5, 3)
+            .sampleSizes(List.of(5, 3))
             .batchSize(5)
             .maxIterations(100)
             .randomSeed(42L)
@@ -276,16 +278,16 @@ void testLossesWithPoolAggregator() {
         assertThat(epochLosses).isInstanceOf(List.class);
         assertThat(((List<Double>) epochLosses).stream().mapToDouble(Double::doubleValue).toArray())
             .contains(new double[]{
-                    90.53,
-                    83.29,
-                    74.75,
-                    74.61,
-                    74.68,
-                    74.54,
-                    74.46,
-                    74.47,
-                    74.41,
-                    74.41
+                87.34,
+                80.75,
+                74.07,
+                93.12,
+                96.36,
+                80.50,
+                77.31,
+                99.70,
+                83.60,
+                83.60
                 }, Offset.offset(0.05)
             );
     }
@@ -306,6 +308,35 @@ void testConvergence() {
         assertThat(trainMetrics.ranIterationsPerEpoch()).containsExactly(2);
     }
 
+    @ParameterizedTest
+    @CsvSource({
+        "0.01, true, 8",
+        "1.0, false, 10"
+    })
+    void batchesPerIteration(double batchSamplingRatio, boolean expectedConvergence, int expectedRanEpochs) {
+        var trainer = new GraphSageModelTrainer(
+            configBuilder.modelName("convergingModel:)")
+                .maybeBatchSamplingRatio(batchSamplingRatio)
+                .embeddingDimension(12)
+                .aggregator(AggregatorType.POOL)
+                .epochs(10)
+                .tolerance(1e-10)
+                .sampleSizes(List.of(5, 3))
+                .batchSize(5)
+                .maxIterations(100)
+                .randomSeed(42L)
+                .build(),
+            Pools.DEFAULT,
+            ProgressTracker.NULL_TRACKER
+        );
+
+        var trainResult = trainer.train(graph, features);
+
+        var trainMetrics = trainResult.metrics();
+        assertThat(trainMetrics.didConverge()).isEqualTo(expectedConvergence);
+        assertThat(trainMetrics.ranEpochs()).isEqualTo(expectedRanEpochs);
+    }
+
     @ParameterizedTest
     @ValueSource(longs = {20L, -100L, 30L})
     void seededSingleBatch(long seed) {