plot_learning_curve.py

import numpy as np
import matplotlib.pyplot as plt
from sklearn import cross_validation
from sklearn.naive_bayes import GaussianNB
from sklearn.svm import SVC
from sklearn.datasets import load_digits
from sklearn.learning_curve import learning_curve

def plot_learning_curve(estimator, title, X, y, ylim=None, cv=None,
                        n_jobs=1, train_sizes=np.linspace(.1,1.0,5)):
    '''
    generate a simple plot of the test and training learning curve
    parameters:
    estimator, estimator
    title is title
    X, y
    ylim is the bounds of y
    cv are the crossvalidation


    '''

    plt.figure()
    plt.title(title)
    if ylim is not None:
        plt.ylim(*ylim)
    plt.xlabel('training examples')
    plt.ylabel("score")
    train_sizes, train_scores, test_scores = learning_curve(estimator, X, y, cv=cv, n_jobs=n_jobs,
                                                            train_sizes=train_sizes)
    train_scores_mean = np.mean(train_scores, axis=1)
    test_scores_mean = np.mean(test_scores, axis=1)
    test_scores_std = np.std(test_scores, axis=1)
    train_scores_std = np.std(train_scores, axis=1)

    plt.grid()

    plt.fill_between(train_sizes, train_scores_mean - train_scores_std,
                     train_scores_mean + train_scores_std, alpha=0.1,
                     color = 'r')
    plt.fill_between(train_sizes, test_scores_mean - test_scores_std,
                     test_scores_mean + test_scores_std, alpha=0.1,
                     color = 'g')
    plt.legend(loc='best')
    return plt

def main():
    digits = load_digits()
    X, y = digits.data, digits.target

    title = 'learning curves (naive bayes)'
    #cross validation with 100 iterations to get smoother mean test and train
    # score curves, each time with 20% data randomly selected as a validation set.

    cv = cross_validation.ShuffleSplit(digits.data.shape[0], n_iter=100, test_size=0.2, random_state=0)

    estimator = GaussianNB()
    plot_learning_curve(estimator, title, X, y, ylim=(0.7, 1.01), cv=cv, n_jobs=4)

    title = 'learning curves (SVM, RBF kernel, $/gamma=0.001%)'
    #SVC is more expensive so we use fewer CV iterations

    cv = cross_validation.ShuffleSplit(digits.data.shape[0], n_iter=10,
                                       test_size=0.2, random_state=0)

    estimator = SVC(gamma=0.001)
    plot_learning_curve(estimator, title, X, y, ylim=(0.7, 1.01), cv=cv, n_jobs=4)
    plt.show()


if __name__ == '__main__':
    main()