(args)
| 59 | |
| 60 | |
| 61 | def generate_train_val_test(args): |
| 62 | df = pd.read_hdf(args.traffic_df_filename) |
| 63 | # 0 is the latest observed sample. |
| 64 | x_offsets = np.sort( |
| 65 | # np.concatenate(([-week_size + 1, -day_size + 1], np.arange(-11, 1, 1))) |
| 66 | np.concatenate((np.arange(-11, 1, 1),)) |
| 67 | ) |
| 68 | # Predict the next one hour |
| 69 | y_offsets = np.sort(np.arange(1, 13, 1)) |
| 70 | # x: (num_samples, input_length, num_nodes, input_dim) |
| 71 | # y: (num_samples, output_length, num_nodes, output_dim) |
| 72 | x, y, index = generate_graph_seq2seq_io_data( |
| 73 | df, |
| 74 | x_offsets=x_offsets, |
| 75 | y_offsets=y_offsets, |
| 76 | add_time_in_day=True, |
| 77 | add_day_in_week=False, |
| 78 | ) |
| 79 | time = df.index |
| 80 | |
| 81 | print("x shape: ", x.shape, ", y shape: ", y.shape) |
| 82 | # Write the data into npz file. |
| 83 | # num_test = 6831, using the last 6831 examples as testing. |
| 84 | # for the rest: 7/8 is used for training, and 1/8 is used for validation. |
| 85 | num_samples = time.shape[0] |
| 86 | num_test = round(num_samples * 0.2) |
| 87 | num_train = round(num_samples * 0.7) |
| 88 | num_val = num_samples - num_test - num_train |
| 89 | |
| 90 | # train |
| 91 | x_train, y_train, time_train = x[:num_train], y[:num_train], index[:num_train] |
| 92 | # val |
| 93 | x_val, y_val, time_val = ( |
| 94 | x[num_train: num_train + num_val], |
| 95 | y[num_train: num_train + num_val], |
| 96 | index[num_train: num_train + num_val], |
| 97 | ) |
| 98 | # test |
| 99 | x_test, y_test, time_test = x[-num_test:], y[-num_test:], index[-num_test:] |
| 100 | |
| 101 | for cat in ["train", "val", "test"]: |
| 102 | _x, _y, _time = locals()["x_" + cat], locals()["y_" + cat], locals()["time_" + cat] |
| 103 | print(cat, "x: ", _x.shape, "y:", _y.shape) |
| 104 | np.savez_compressed( |
| 105 | os.path.join(args.output_dir, "%s.npz" % cat), |
| 106 | x=_x, |
| 107 | y=_y, |
| 108 | time = _time, |
| 109 | x_offsets=x_offsets.reshape(list(x_offsets.shape) + [1]), |
| 110 | y_offsets=y_offsets.reshape(list(y_offsets.shape) + [1]), |
| 111 | time_offsets=y_offsets.reshape(list(y_offsets.shape) + [1]), |
| 112 | ) |
| 113 | |
| 114 | |
| 115 | def main(args): |
no test coverage detected