首先，定义一个数据获取的函数，其中通过baostock库获取指定代码股票从2020/1/1到2022/1/1两年的日线数据。
def get_stockdata(code): rs = bs.query_history_k_data_plus(code, 'date,close,volume,turn', start_date='2020-01-01', end_date='2022-01-01',
frequency='d', adjustflag='2')
return rs.get_data()
之后，定义预测长度以及需要获取的股票代码。然后将获取到的数据划分为训练集和测试集。这里prediction_length变量定义了预测长度，stock_list中选取了三支上海证券交易所中三支医疗板块的股票。之后，将得到的股票数据转换为GluonTS指定的输入数据格式ListDataset。其中，start表示预测起始值，target表示预测的目标变量，cat表示引入的静态变量，这里用到了股票的id，dynamic_feat表示其他的动态标量，这里引入了交易量跟换手率，它的长度需要跟target一致。prediction_length = 20stock_list = ['sh.600227', 'sh.600200', 'sh.600201']train_dic_list = []test_dic_list = []lg = bs.login()
for stock_id in stock_list: df = get_stockdata(stock_id) train_dic = {'start':df.date[0], 'target':df.close, 'cat':int(stock_id.split('.')[1]), 'dynamic_feat':[df.volume, df.turn]} test_dic = {'start':df.date[0], 'target':df.close[:-prediction_length], 'cat':int(stock_id.split('.')[1]), 'dynamic_feat':[df.volume[:-prediction_length], df.turn[:-prediction_length]]} train_dic_list.append(train_dic) test_dic_list.append(test_dic)
bs.logout()
之后，需要构造DeepAR模型，并进行训练。其中，prediction_length：预测范围的长度；context_length表示在计算预测之前要为RNN展开的步骤数（默认context_length等于prediction_length）；num_layers表示RNN层数；num_cells表示每层的RNN的神经元个数。estimator = DeepAREstimator( prediction_length=prediction_length, context_length=60, freq='1d', num_layers=2, num_cells=64, trainer=Trainer(epochs=20, learning_rate=1e-2, num_batches_per_epoch=32))predictor = estimator.train(train_data)
def plot_prob_forecasts(ts_entry, forecast_entry, path, sample_id): plot_length = 150 prediction_intervals = (50, 80) legend = ['observations', 'median prediction'] + [f'{k}% prediction interval' for k in prediction_intervals][::-1]
_, ax = plt.subplots(1, 1, figsize=(10, 7)) ts_entry[-plot_length:].plot(ax=ax) forecast_entry.plot(prediction_intervals=prediction_intervals, color='g') ax.axvline(ts_entry.index[-prediction_length], color='r') plt.legend(legend, loc='upper left') plt.savefig('{}forecast_{}.png'.format(path, sample_id))
plt.close()
之后调用模型评估的方法对训练好的模型进行评估，其中num_samples表示可视化数据的长度。最后，将plot的结果保存在本地。
forecast_it, ts_it = make_evaluation_predictions( dataset=test_data, predictor=predictor, num_samples=100)
tss = list(tqdm(ts_it, total=len(test_data)))forecasts = list(tqdm(forecast_it, total=len(test_data)))
plot_log_path = './plots/'directory = os.path.dirname(plot_log_path)if not os.path.exists(directory): os.makedirs(directory) for i in tqdm(range(len(stock_list))): ts_entry = tss[i] forecast_entry = forecasts[i] plot_prob_forecasts(ts_entry, forecast_entry, plot_log_path, i)
经过20个epoch的训练，可以看出模型的loss不断下降，实际应用中可以通过超参优化，warmup，early stop，dropout等一系列的操作来使得模型具有更好的训练效果，这里不再赘述。
100%|██████████| 32/32 [00:12<00:00, 2.63it/s, epoch=1/20, avg_epoch_loss=1.99]100%|██████████| 32/32 [00:13<00:00, 2.34it/s, epoch=2/20, avg_epoch_loss=1.01]100%|██████████| 32/32 [00:07<00:00, 4.49it/s, epoch=3/20, avg_epoch_loss=1.07]100%|██████████| 32/32 [00:06<00:00, 4.61it/s, epoch=4/20, avg_epoch_loss=0.87]100%|██████████| 32/32 [00:06<00:00, 4.63it/s, epoch=5/20, avg_epoch_loss=0.614]100%|██████████| 32/32 [00:05<00:00, 5.40it/s, epoch=6/20, avg_epoch_loss=0.595]100%|██████████| 32/32 [00:05<00:00, 6.18it/s, epoch=7/20, avg_epoch_loss=0.474]100%|██████████| 32/32 [00:05<00:00, 5.82it/s, epoch=8/20, avg_epoch_loss=0.399]100%|██████████| 32/32 [00:05<00:00, 5.85it/s, epoch=9/20, avg_epoch_loss=0.36]100%|██████████| 32/32 [00:05<00:00, 5.58it/s, epoch=10/20, avg_epoch_loss=0.273]100%|██████████| 32/32 [00:05<00:00, 5.78it/s, epoch=11/20, avg_epoch_loss=0.208]100%|██████████| 32/32 [00:05<00:00, 5.83it/s, epoch=12/20, avg_epoch_loss=0.169]100%|██████████| 32/32 [00:05<00:00, 6.04it/s, epoch=13/20, avg_epoch_loss=0.156]100%|██████████| 32/32 [00:05<00:00, 6.04it/s, epoch=14/20, avg_epoch_loss=0.26]100%|██████████| 32/32 [00:05<00:00, 5.98it/s, epoch=15/20, avg_epoch_loss=0.189]100%|██████████| 32/32 [00:05<00:00, 5.69it/s, epoch=16/20, avg_epoch_loss=0.138]100%|██████████| 32/32 [00:05<00:00, 5.62it/s, epoch=17/20, avg_epoch_loss=0.0314]100%|██████████| 32/32 [00:05<00:00, 5.79it/s, epoch=18/20, avg_epoch_loss=-.033]100%|██████████| 32/32 [00:05<00:00, 5.73it/s, epoch=19/20, avg_epoch_loss=0.0111]100%|██████████| 32/32 [00:05<00:00, 5.67it/s, epoch=20/20, avg_epoch_loss=0.00313]
三只股票的预测结果可视化如下三张图所示，其中蓝色的线代表ground truth，绿色的线代表概率预测的中值，两个绿色区域则代表了80%跟50%的置信区间。从可视化的结果中也可以看出，随着时间的推移，模型预测结果的误差逐渐扩大，这也侧面说明了进行多步预测的难度之大。
