objectImageIdentifierAI/nnTrain/train.py
nnTrain/train.pyRaw
1import matplotlib.pyplot as plt
2import numpy as np
3from random import randint
4import pickle
5from sknn.mlp import Classifier, Layer, Convolution
6
7def datasetToTrainAndTestData(dataset, numtest):
8 np.random.shuffle(dataset)
9 print "length total data:" + str(len(dataset))
10
11 traindata = np.copy(dataset)
12 testdata = []
13 for i in range(numtest):
14 #get random integer between 0 and the total amount of images in the dataset
15 n = randint(0, len(traindata))
16 testdata.append(dataset[n])
17
18 #delete the n image (dataset[n]) of the traindata
19 traindata = np.delete(traindata, n, axis=0)
20 testdataNP = np.array(testdata)
21 return traindata, testdataNP
22
23
24#read the dataset made with the 'imagesToDataset' repository
25dataset = np.load('dataset.npy')
26
27traindata, testdata = datasetToTrainAndTestData(dataset, 10)
28print "length traindata: " + str(len(traindata))
29print "length testdata: " + str(len(testdata))
30
31#traindataAttributes contains all the pixels of each image
32traindataAttributes = traindata[:,0]
33traindataAttributes = np.array([[row] for row in traindataAttributes])
34
35#traindataLabels contains each label of each image
36traindataLabels = traindata[:,1]
37traindataLabels = traindataLabels.astype('int')
38
39#testdataAttributes contains the pixels of the test images
40testdataAttributes = testdata[:,0]
41testdataAttributes = np.array([[row] for row in testdataAttributes])
42
43#testdataLabels contains each label of each image
44testdataLabels = testdata[:,1]
45testdataLabels = testdataLabels.astype('int')
46
47#default: units=100, learning_rate=0.001, n_iter=25
48nn = Classifier(
49 layers=[
50 Layer("Sigmoid", units=10),
51 Layer("Softmax")],
52 learning_rate=0.001,
53 n_iter=20,
54 verbose=True)
55
56nn.fit(traindataAttributes, traindataLabels)
57
58print('\nTRAIN SCORE', nn.score(traindataAttributes, traindataLabels))
59print('TEST SCORE', nn.score(testdataAttributes, testdataLabels))
60
61#save the neural network configuration
62pickle.dump(nn, open('nn.pkl', 'wb'))