-
Notifications
You must be signed in to change notification settings - Fork 1
/
pretrain_main.py
79 lines (62 loc) · 3.61 KB
/
pretrain_main.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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import os
import numpy as np
import pandas as pd
import torch
from sklearn.metrics import r2_score
from models.common_net import M100_NN, M50_NN
from torch.utils.data import Dataset, DataLoader
from models.customer_dataset import PretrainDataset
from utils.training_utils import BaseLearner, forecast, train
from models.customer_dataset import PretrainDataset
# project info
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
save_path = os.path.join(os.getcwd(), 'checkpoint')
basin_id = 6431500
# load data
# train data
train_data_df = pd.read_csv(r'data/{}/train_data_df.csv'.format(basin_id))
# test data
test_data_df = pd.read_csv(r'data/{}/test_data_df.csv'.format(basin_id))
means = train_data_df[['S_snow', 'S_water', 'Precp', 'Temp', 'Lday']].mean().values
stds = train_data_df[['S_snow', 'S_water', 'Precp', 'Temp', 'Lday']].std().values
ET_train_dataloader = DataLoader(dataset=PretrainDataset(
train_data_df, input_cols=['S_snow', 'S_water', 'Temp'], target_cols=['ET_mech']),
batch_size=32, shuffle=True)
ET_test_dataloader = DataLoader(dataset=PretrainDataset(
test_data_df, input_cols=['S_snow', 'S_water', 'Temp'], target_cols=['ET_mech']),
batch_size=32, shuffle=False)
Q_train_dataloader = DataLoader(dataset=PretrainDataset(
train_data_df, input_cols=['S_water', 'Precp'], target_cols=['Q_mech']),
batch_size=32, shuffle=False)
Q_test_dataloader = DataLoader(dataset=PretrainDataset(
test_data_df, input_cols=['S_water', 'Precp'], target_cols=['Q_mech']),
batch_size=32, shuffle=False)
net_train_dataloader = DataLoader(dataset=PretrainDataset(
train_data_df, input_cols=['S_snow', 'S_water', 'Precp', 'Temp'],
target_cols=['ET_mech', 'Q_mech', 'M_mech', 'Ps_mech', 'Pr_mech']),
batch_size=32, shuffle=False)
net_test_dataloader = DataLoader(dataset=PretrainDataset(
test_data_df, input_cols=['S_snow', 'S_water', 'Precp', 'Temp'],
target_cols=['ET_mech', 'Q_mech', 'M_mech', 'Ps_mech', 'Pr_mech']),
batch_size=32, shuffle=False)
ET_model = M50_NN(input_dim=3, output_dim=1, hidden_units=16, means=means[[0, 1, 3]], stds=stds[[0, 1, 3]]).to(device)
ET_learner = BaseLearner(ET_model, loss_metric=torch.nn.MSELoss())
ET_pretrain_save_path = os.path.join(save_path, str(basin_id), 'pretrain', 'M50-ET')
ET_trained_model, trained_ET_learner = train(ET_learner, ET_train_dataloader, ET_pretrain_save_path, save_pt=True)
et_real_arr, et_pred_arr = forecast(trained_ET_learner, ET_test_dataloader)
print('test acc ' + str(r2_score(et_real_arr, et_pred_arr)))
Q_model = M50_NN(input_dim=2, output_dim=1, hidden_units=16, means=means[[1, 2]], stds=stds[[1, 2]]).to(device)
Q_learner = BaseLearner(Q_model, loss_metric=torch.nn.MSELoss())
Q_pretrain_save_path = os.path.join(save_path, str(basin_id), 'pretrain', 'M50-Q')
Q_trained_model, trained_Q_learner = train(Q_learner, Q_train_dataloader, Q_pretrain_save_path, save_pt=True)
net_model = M100_NN(input_dim=4, output_dim=5, hidden_units=32,
means=means[[0, 1, 2, 3]], stds=stds[[0, 1, 2, 3]]).to(device)
net_learner = BaseLearner(net_model, loss_metric=torch.nn.MSELoss(), lr=1e-3)
net_pretrain_save_path = os.path.join(save_path, str(basin_id), 'pretrain', 'M100')
net_trained_model, trained_net_learner = train(net_learner, net_train_dataloader, net_pretrain_save_path, save_pt=True)
# q_real_arr, q_pred_arr = forecast(trained_Q_learner, Q_test_dataloader)
# q_real_arr, q_pred_arr = forecast(trained_Q_learner, Q_train_dataloader)
# Q_obs = train_data_df['Q_obs'].values
# q_pred = np.exp(q_pred_arr)
# print('test acc ' + str(r2_score(q_real_arr, q_pred_arr)))
# print('test acc ' + str(r2_score(Q_obs, q_pred)))