Add functionality to restart tree training where previously left off.

This commit is contained in:
Joel Therrien 2018-08-09 16:34:10 -07:00
parent 6d65d48844
commit e92abdab13

View file

@ -8,6 +8,7 @@ import lombok.AccessLevel;
import lombok.AllArgsConstructor; import lombok.AllArgsConstructor;
import lombok.Builder; import lombok.Builder;
import java.io.File;
import java.io.FileOutputStream; import java.io.FileOutputStream;
import java.io.IOException; import java.io.IOException;
import java.io.ObjectOutputStream; import java.io.ObjectOutputStream;
@ -122,10 +123,19 @@ public class ForestTrainer<Y, TO, FO> {
public void trainParallelOnDisk(int threads){ public void trainParallelOnDisk(int threads){
final ExecutorService executorService = Executors.newFixedThreadPool(threads); // First we need to see how many trees there currently are
final AtomicInteger treeCount = new AtomicInteger(0); // tracks how many trees are finished final File folder = new File(saveTreeLocation);
if(!folder.isDirectory()){
throw new IllegalArgumentException("Tree directory must be a directory!");
}
for(int j=0; j<ntree; j++){
final File[] treeFiles = folder.listFiles(((file, s) -> s.endsWith(".tree")));
final ExecutorService executorService = Executors.newFixedThreadPool(threads);
final AtomicInteger treeCount = new AtomicInteger(treeFiles.length); // tracks how many trees are finished
for(int j=treeCount.get(); j<ntree; j++){
final Runnable worker = new TreeSavedWorker(data, "tree-" + (j+1) + ".tree", treeCount); final Runnable worker = new TreeSavedWorker(data, "tree-" + (j+1) + ".tree", treeCount);
executorService.execute(worker); executorService.execute(worker);
} }