package cc.mallet.fst;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.logging.Logger;
import cc.mallet.types.InstanceList;
import cc.mallet.types.MatrixOps;
import cc.mallet.optimize.Optimizable;
import cc.mallet.util.MalletLogger;
/**
* An adaptor for optimizables based on batch values/gradients.
* <p>
* Computes values, gradients for each batch in multiple threads and combines
* them in the end.
*
* @author Gaurav Chandalia
* @see CRFOptimizableByBatchLabelLikelihood
*/
public class ThreadedOptimizable implements Optimizable.ByGradientValue {
private static Logger logger = MalletLogger.getLogger(ThreadedOptimizable.class.getName());
/** Data */
protected InstanceList trainingSet;
/** optimizable to be parallelized */
protected Optimizable.ByCombiningBatchGradient optimizable;
/** Value obtained from the optimizable for each batch */
protected double[] batchCachedValue;
/** Gradient obtained from the optimizable for each batch */
protected List<double[]> batchCachedGradient;
// determine when value/gradient become stale
protected CacheStaleIndicator cacheIndicator;
// tasks to be executed in individual threads, each task is instantiated only
// once but executed in every iteration
private transient Collection<Callable<Double>> valueTasks;
private transient Collection<Callable<Boolean>> gradientTasks;
// thread pool to compute value/gradient for one batch of data
private transient ThreadPoolExecutor executor;
// milliseconds
public static final int SLEEP_TIME = 100;
/**
* Initializes the optimizable and starts new threads.
*
* @param optimizable Optimizable to be parallelized
* @param numFactors Number of factors in model's parameters, used to
* initialize the gradient
* @param cacheIndicator Determines when value/gradient become stale
*/
public ThreadedOptimizable(Optimizable.ByCombiningBatchGradient optimizable,
InstanceList trainingSet, int numFactors,
CacheStaleIndicator cacheIndicator) {
// set up
this.trainingSet = trainingSet;
this.optimizable = optimizable;
int numBatches = optimizable.getNumBatches();
assert(numBatches > 0) : "Invalid number of batches: " + numBatches;
batchCachedValue = new double[numBatches];
batchCachedGradient = new ArrayList<double[]>(numBatches);
for (int i = 0; i < numBatches; ++i) {
batchCachedGradient.add(new double[numFactors]);
}
this.cacheIndicator = cacheIndicator;
logger.info("Creating " + numBatches + " threads for updating gradient...");
executor = (ThreadPoolExecutor) Executors.newFixedThreadPool(numBatches);
this.createTasks();
}
public Optimizable.ByCombiningBatchGradient getOptimizable() {
return optimizable;
}
/**
* Shuts down the executor used to start and run threads to compute values
* and gradients.
* <p>
* *Note*: For a clean exit of all the threads, it is recommended to call
* this method after training finishes.
*/
public void shutdown() {
// fix submitted by Mark Dredze (mdredze@cs.jhu.edu)
executor.shutdown();
try {
executor.awaitTermination(30, TimeUnit.SECONDS);
} catch (InterruptedException e) {
e.printStackTrace();
}
assert(executor.shutdownNow().size() == 0) : "All tasks didn't finish";
}
public double getValue () {
if (cacheIndicator.isValueStale()) {
// compute values again
try {
// run all threads and wait for them to finish
List<Future<Double>> results = executor.invokeAll(valueTasks);
// compute final log probability
int batch = 0;
for (Future<Double> f : results) {
try {
batchCachedValue[batch++] = f.get();
} catch (ExecutionException ee) {
ee.printStackTrace();
}
}
} catch (InterruptedException ie) {
ie.printStackTrace();
}
double cachedValue = MatrixOps.sum(batchCachedValue);
logger.info("getValue() (loglikelihood, optimizable by label likelihood) =" + cachedValue);
return cachedValue;
}
return MatrixOps.sum(batchCachedValue);
}
/**
* Returns the gradient, re-computes if gradient is stale. <p>
*
* *Note*: Assumes that <tt>buffer</tt> is already initialized.
*/
public void getValueGradient (double[] buffer) {
if (cacheIndicator.isGradientStale()) {
// compute values again if required
this.getValue();
// compute gradients again
try {
// run all threads and wait for them to finish
executor.invokeAll(gradientTasks);
} catch (InterruptedException ie) {
ie.printStackTrace();
}
}
optimizable.combineGradients(batchCachedGradient, buffer);
}
/**
* Creates tasks to be executed in parallel, each task looks at a batch of
* data.
*/
protected void createTasks() {
int numBatches = optimizable.getNumBatches();
valueTasks = new ArrayList<Callable<Double>>(numBatches);
gradientTasks = new ArrayList<Callable<Boolean>>(numBatches);
// number of instances per batch
int numBatchInstances = trainingSet.size() / numBatches;
// batch assignments
int start = -1, end = -1;
for (int i = 0; i < numBatches; ++i) {
// get the indices of batch
if (i == 0) {
start = 0;
end = start + numBatchInstances;
} else if (i == numBatches-1) {
start = end;
end = trainingSet.size();
} else {
start = end;
end = start + numBatchInstances;
}
valueTasks.add(new ValueHandler(i, new int[]{start, end}));
gradientTasks.add(new GradientHandler(i, new int[]{start, end}));
}
}
public int getNumParameters () { return optimizable.getNumParameters(); }
public void getParameters (double[] buffer) {
optimizable.getParameters(buffer);
}
public double getParameter (int index) {
return optimizable.getParameter(index);
}
public void setParameters (double [] buff) {
optimizable.setParameters(buff);
}
public void setParameter (int index, double value) {
optimizable.setParameter(index, value);
}
/**
* Computes value in a separate thread for a batch of data.
*/
private class ValueHandler implements Callable<Double> {
private int batchIndex;
private int[] batchAssignments;
public ValueHandler(int batchIndex, int[] batchAssignments) {
this.batchIndex = batchIndex;
this.batchAssignments = batchAssignments;
}
/**
* Returns the value for a batch.
*/
public Double call() {
return optimizable.getBatchValue(batchIndex, batchAssignments);
}
}
/**
* Computes gradient in a separate thread for a batch of data.
*/
private class GradientHandler implements Callable<Boolean> {
private int batchIndex;
private int[] batchAssignments;
public GradientHandler(int batchIndex, int[] batchAssignments) {
this.batchIndex = batchIndex;
this.batchAssignments = batchAssignments;
}
/**
* Computes the gradient for a batch, always returns true.
*/
public Boolean call() {
optimizable.getBatchValueGradient(batchCachedGradient.get(batchIndex),
batchIndex, batchAssignments);
return true;
}
}
}