Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 23 additions & 3 deletions src/Microsoft.ML.Vision/ImageClassificationTrainer.cs
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,8 @@ namespace Microsoft.ML.Vision
///
/// ### Training Algorithm Details
/// Trains a Deep Neural Network(DNN) by leveraging an existing pre-trained model such as Resnet50 for the purpose
/// of classifying images.
/// of classifying images. The technique was inspired from [TensorFlow's retrain image classification tutorial]
/// (https://www.tensorflow.org/hub/tutorials/image_retraining)
/// ]]>
/// </format>
/// </remarks>
Expand Down Expand Up @@ -392,7 +393,7 @@ public sealed class Options : TrainerInputBaseWithLabel
public Action<ImageClassificationMetrics> MetricsCallback = null;

/// <summary>
/// Indicates the path where the models get downloaded to and cache files saved, default is a new temporary directory
/// Indicates the path where the image bottleneck cache files and trained model are saved, default is a new temporary directory
/// </summary>
[Argument(ArgumentType.AtMostOnce, HelpText = "Indicates the path where the models get downloaded to and cache files saved, default is a new temporary directory.", SortOrder = 15)]
public string WorkspacePath = null;
Expand Down Expand Up @@ -591,6 +592,7 @@ private void InitializeTrainingGraph(IDataView input)
_classCount = labelCount == 1 ? 2 : (int)labelCount;
var imageSize = ImagePreprocessingSize[_options.Arch];
_session = LoadTensorFlowSessionFromMetaGraph(Host, _options.Arch).Session;
_session.graph.as_default();
(_jpegData, _resizedImage) = AddJpegDecoding(imageSize.Item1, imageSize.Item2, 3);
_jpegDataTensorName = _jpegData.name;
_resizedImageTensorName = _resizedImage.name;
Expand Down Expand Up @@ -631,6 +633,14 @@ private protected override MulticlassPredictionTransformer<ImageClassificationMo

private protected override ImageClassificationModelParameters TrainModelCore(TrainContext trainContext)
{
// Workspace directory is cleaned after training run. However, the pipeline can be re-used by calling
// fit() again after transform(), in which case we must ensure workspace directory exists. This scenario
// is typical in the case of cross-validation.
if (!Directory.Exists(_options.WorkspacePath))
{
Directory.CreateDirectory(_options.WorkspacePath);
}

InitializeTrainingGraph(trainContext.TrainingSet.Data);
CheckTrainingParameters(_options);
var validationSet = trainContext.ValidationSet?.Data ?? _options.ValidationSet;
Expand Down Expand Up @@ -1301,7 +1311,7 @@ private void VariableSummaries(RefVariable var)
var optimizer = useLearningRateScheduler ? tf.train.GradientDescentOptimizer(_learningRateInput) :
tf.train.GradientDescentOptimizer(learningRate);

_trainStep = optimizer.minimize(crossEntropyMean);
_trainStep = optimizer.minimize(crossEntropyMean);
});

return (_trainStep, crossEntropyMean, _labelTensor, _softMaxTensor);
Expand Down Expand Up @@ -1341,6 +1351,11 @@ private void Dispose(bool disposing)
{
_session.close();
}

if (_session != null && _session.graph != IntPtr.Zero)
{
_session.graph.Dispose();
}
}

/// <summary>
Expand Down Expand Up @@ -1527,6 +1542,11 @@ private void Dispose(bool disposing)
{
_session.close();
}

if (_session != null && _session.graph != IntPtr.Zero)
{
_session.graph.Dispose();
}
}
}
}