forked from wangjin0818/Kaggle-MNIST-Competition
-
Notifications
You must be signed in to change notification settings - Fork 0
/
naive_bayes.py
30 lines (23 loc) · 903 Bytes
/
naive_bayes.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
from sklearn.naive_bayes import GaussianNB
from sklearn.naive_bayes import MultinomialNB
import numpy as np
import util
def GaussianNBClassify(trainData, trainLabel, testData):
nbClf = GaussianNB()
nbClf.fit(trainData, np.ravel(trainLabel))
testLabel = nbClf.predict(testData)
util.saveResult(testLabel, './result/gaussion_nb_result.csv')
return testLabel
def MultinomialNBClassify(trainData, trainLabel, testData):
nbClf = MultinomialNB()
nbClf.fit(trainData, np.ravel(trainLabel))
testLabel = nbClf.predict(testData)
util.saveResult(testLabel, './result/multinomial_nb_result.csv')
return testLabel
if __name__ == "__main__":
trainData, trainLabel = util.loadTrainData()
testData = util.loadTestData()
GaussianNBClassify(trainData, trainLabel, testData)
print 'gaussion_nb is finished!'
MultinomialNBClassify(trainData, trainLabel, testData)
print 'multinomial_nb is finished!'