SGST_stanford_main.py

'''
author:Donglin Zhou
time:2021.12.15
function:基于市场情绪和技术面分析的股票价格预测
使用的情感数据集为stanfordNLP
'''
import os
import time
import argparse
from att_utils import *
from SGST_stanford_model import SGST
import torch
from torch import nn

class MSE_loss(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, y):
        return torch.mean(torch.pow((x - y), 2))

def str2bool(s):
    if s not in {'false', 'true'}:
        raise ValueError('Not a valid boolean string')
    return s == 'true'

if __name__ == '__main__':
    # 读取数据的路径
    datafile_name = 'ST_MSFT_data_PCA.csv'
    file_name = './dataset/stockstats&stanford/' + datafile_name
    datatype = 'stanford'

    parser = argparse.ArgumentParser()
    parser.add_argument('--sparse_rate', default=3, type=int)       # 稀疏率，1,3,5
    parser.add_argument('--lr', default=0.001, type=float)
    parser.add_argument('--look_back', default=5, type=int)        # 5 10 15 20 25
    parser.add_argument('--hidden_units', default=50, type=int)      # 5 10 15 20 50
    parser.add_argument('--num_epochs', default=200, type=int)
    parser.add_argument('--batch_size', default=8, type=int)        # 5
    parser.add_argument('--num_heads', default=1, type=int)
    parser.add_argument('--num_blocks', default=2, type=int)        # 1 2 4
    parser.add_argument('--dropout_rate', default=0.1, type=float)
    parser.add_argument('--test_size', default=0.3, type=float)
    parser.add_argument('--device', default='cuda', type=str)
    args = parser.parse_args()
    # 切分数据集
    X_train, y_train, X_test, y_test = get_data_att(file_name, datatype, args.look_back, args.test_size)

    model = SGST(args).to(args.device)
    num_batch = int(len(X_train) / args.batch_size)
    # 参数初始化
    for name, param in model.named_parameters():
        try:
            # 使得初始参数服从均匀分布
            torch.nn.init.xavier_uniform_(param.data)
        except:
            pass

    model.train()
    epoch_start_idx = 1

    criterion = MSE_loss()                                                                  # MSE
    adam_optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, betas=(0.9, 0.98))    # 亚当优化器

    T = 0.0
    t0 = time.time()
    all_data = []       # 存放训练结果

    # 进入迭代训练
    for epoch in range(epoch_start_idx, args.num_epochs + 1):
        for step in range(num_batch):
            # 获取数据并转化维度
            feature1 = X_train[step:step+args.batch_size, :, 0]
            feature2 = X_train[step:step+args.batch_size, :, 1]
            feature3 = X_train[step:step+args.batch_size, :, 2]
            v_Positive = X_train[step:step + args.batch_size, :, 3]
            Positive = X_train[step:step+args.batch_size, :, 4]
            confidence = X_train[step:step + args.batch_size, :, 5]
            Negative = X_train[step:step+args.batch_size, :, 6]
            v_Negative = X_train[step:step + args.batch_size, :, 7]
            y_seq = y_train[step:step+args.batch_size,:, :]

            feature1 = feature1.reshape(feature1.shape[0], feature1.shape[1], 1)
            feature2 = feature2.reshape(feature2.shape[0], feature2.shape[1], 1)
            feature3 = feature3.reshape(feature3.shape[0], feature3.shape[1], 1)
            Positive = Positive.reshape(Positive.shape[0], Positive.shape[1], 1)
            Negative = Negative.reshape(Negative.shape[0], Negative.shape[1], 1)
            v_Positive = v_Positive.reshape(v_Positive.shape[0], v_Positive.shape[1], 1)
            confidence = confidence.reshape(confidence.shape[0], confidence.shape[1], 1)
            v_Negative = v_Negative.reshape(v_Negative.shape[0], v_Negative.shape[1], 1)

            pos_logits, seqs = model(y_seq, feature1, feature2, feature3, v_Positive, Positive,confidence, Negative, v_Negative)

            adam_optimizer.zero_grad()
            loss = criterion(pos_logits, seqs)
            loss.backward()
            adam_optimizer.step()

            print("loss in epoch {} iteration {}: {}".format(epoch, step, loss.item()))

        if epoch % 20 == 0:
            data_list = []
            model.eval()
            t1 = time.time() - t0
            T += t1
            print('Evaluating', end='')
            mse, mape, smape = evaluate_test(model, X_test, y_test, datatype)
            print('epoch:%d, time: %f(s),  test (mse: %.4f, mape:%.4f, smape:%.4f)'% (epoch, T, mse, mape,smape))

            data_list.append(float(mse))
            data_list.append(float(mape))
            data_list.append(float(smape))
            all_data.append(data_list)
            t0 = time.time()
            model.train()

        # 训练完成
        if epoch == args.num_epochs:
            # 结果保存的路径
            savepath = './result/SGST/' + datatype + '/' + datafile_name.split('.')[0] + '/'
            if not os.path.isdir(savepath):
                os.makedirs(savepath)
            fname = savepath + 'epoch={}.lr={}.layer={}.head={}.hidden={}.lookback={}.sparserate={}.batchsize={}.test_size={}'
            fname = fname.format(args.num_epochs, args.lr, args.num_blocks, args.num_heads, args.hidden_units,
                                 args.look_back, args.sparse_rate, args.batch_size, args.test_size)
            all_data = pd.DataFrame(all_data)
            all_data.to_csv(fname + '_mse_result.csv', index=None, header=['mse', 'mape', 'smape'])

    print("Done")