99 lines
3.0 KiB
Python
99 lines
3.0 KiB
Python
import os
|
|
import keras
|
|
import numpy as np
|
|
from numpy import zeros, newaxis
|
|
import tensorflow as tf
|
|
from stock.util.Stock2Vector import Stock2Vector
|
|
from classification_models.keras import Classifiers
|
|
|
|
class StockTrainer:
|
|
|
|
RESOURCE_PATH = None
|
|
stock2Vector = None
|
|
|
|
def __init__(self, RESOURCE_PATH):
|
|
self.RESOURCE_PATH = RESOURCE_PATH
|
|
self.stock2Vector = Stock2Vector(RESOURCE_PATH)
|
|
return
|
|
|
|
def getDataset(self, stock_code):
|
|
VECTOR_SIZE = 299
|
|
df, minmax_df = self.stock2Vector.makeTrainData(stock_code)
|
|
|
|
TOTAL_X, TOTAL_Y = [], []
|
|
for key in df:
|
|
if key == "date":
|
|
continue
|
|
elif key == "label":
|
|
TOTAL_Y.append(df[key].tolist())
|
|
else:
|
|
TOTAL_X.append(df[key].tolist())
|
|
|
|
SIZE_WIDTH = len(TOTAL_X[0])
|
|
SIZE_HEIGHT = len(TOTAL_X)
|
|
X, Y = [], []
|
|
for i in range(VECTOR_SIZE, SIZE_WIDTH):
|
|
temp_X, temp_Y = np.zeros((VECTOR_SIZE, VECTOR_SIZE)), np.zeros(0)
|
|
for j in range(SIZE_HEIGHT):
|
|
temp_X[j][0:VECTOR_SIZE] = TOTAL_X[j][i-VECTOR_SIZE:i]
|
|
temp_X = np.stack([temp_X, temp_X, temp_X], axis=-1)
|
|
X.append(temp_X)
|
|
if int(TOTAL_Y[0][i]) == 0:
|
|
Y.append([1, 0, 0])
|
|
elif int(TOTAL_Y[0][i]) == 0.5:
|
|
Y.append([0, 1, 0])
|
|
else:
|
|
Y.append([0, 0, 1])
|
|
if i >= VECTOR_SIZE+10:
|
|
break
|
|
|
|
X = np.asarray(X)
|
|
Y = np.asarray(Y)
|
|
return X, Y
|
|
|
|
def train(self, stock_code):
|
|
ResNet18, preprocess_input = Classifiers.get('inceptionresnetv2')
|
|
|
|
X, Y = self.getDataset(stock_code)
|
|
|
|
X = preprocess_input(X)
|
|
|
|
n_classes = 3
|
|
|
|
# build model
|
|
base_model = ResNet18(input_shape=(299, 299, 3), include_top=False)
|
|
x = tf.keras.layers.GlobalAveragePooling2D()(base_model.output)
|
|
output = tf.keras.layers.Dense(n_classes, activation='softmax')(x)
|
|
model = keras.models.Model(inputs=[base_model.input], outputs=[output])
|
|
|
|
# train
|
|
model.compile(optimizer='SGD', loss='categorical_crossentropy', metrics=['accuracy'])
|
|
|
|
checkpoint_filename = os.path.join(self.RESOURCE_PATH, "model", "stock.ckpt")
|
|
chekpoint = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_filename, save_weights_only=True, verbose=1)
|
|
earlystop = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=3, mode="auto")
|
|
|
|
model.fit(x=X,
|
|
y=Y,
|
|
epochs=10,
|
|
callbacks=[chekpoint, earlystop])
|
|
|
|
return
|
|
|
|
|
|
if __name__ == "__main__":
|
|
PROJECT_HOME = os.path.join(os.path.dirname(__file__))
|
|
RESOURCE_PATH = os.path.join(PROJECT_HOME, "resources")
|
|
|
|
stock_codes = {
|
|
# 252670
|
|
# 122630
|
|
"252670": ['20220729'],
|
|
}
|
|
|
|
for stock_code in stock_codes:
|
|
stockTrainer = StockTrainer(RESOURCE_PATH)
|
|
stockTrainer.train(stock_code)
|
|
|
|
|
|
print ("done...") |