-
Notifications
You must be signed in to change notification settings - Fork 2
/
sde_simple.py
170 lines (148 loc) · 7.49 KB
/
sde_simple.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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
import os
import json
from datetime import datetime
import fire
import torch
from torch.optim import Adam
from models.sde_no_embed import Generator, DiscriminatorSimple, SampledInitialCondition, ColumnwiseInitialCondition, RandomInitialCondition
from models.forcing import SolarMultForcing, WindMultForcing, DawnDuskSampler, SolarAddForcing
from training import WGANGPTrainer
from dataloaders import get_sde_dataloader
from utils.plotting import SDETrainingPlotter
from utils import get_accelerator_device
def train_sdegan(params_file: str = "",
warm_start: bool = False,
epochs: int = 100,
device: str = "cpu",
no_save: bool = False) -> None:
"""
Sets up and trains an SDE-GAN.
Parameters
----------
params_file : str, optional
Path to a JSON file containing the parameters for the model. If None, the parameter set defined
in the function will be used.
warm_start : bool, optional
If True, load saved models and continue training. Requires params_file to be specified.
If False, train a new model.
epochs : int, optional
The number of epochs to train for.
"""
# Set up training. All of these parameters are saved along with the models so the training can be reproduced.
if params_file:
with open(params_file, 'r') as f:
params = json.load(f)
else:
params = {
'time_series_length': 24, # number of nodes in generator output, discriminator input
'ISO': 'ERCOT',
'variables': ['TOTALLOAD', 'WIND', 'SOLAR'],
'time_features': ['HOD'],
'gen_mlp_size': 256,
'gen_num_layers': 3,
'dis_mlp_size': 64,
'dis_num_layers': 1,
'critic_iterations': 10,
'batch_size': 1826,
'gen_lr': 1e-4,
'dis_lr': 1e-4,
'gen_betas': (0.5, 0.9),
'dis_betas': (0.5, 0.9),
'gp_weight': 10.0,
'epochs': epochs,
'total_epochs_trained': 0,
'random_seed': 12345
}
# Find the most appropriate device for training
device = device or get_accelerator_device()
if isinstance(params['variables'], str):
params['variables'] = [params['variables']]
# seed for reproducibility
torch.manual_seed(params['random_seed'])
dataloader, transformer = get_sde_dataloader(iso=params['ISO'],
varname=params['variables'],
segment_size=params['time_series_length'],
time_features=params['time_features'],
batch_size=params['batch_size'],
device=device)
solar = dataloader.dataset[..., params["variables"].index("SOLAR") + 1]
# Some of the data is bad! Solar values must be 0 at night.
bad_rows, _ = torch.where(solar[:, [0, 1, 2, 21, 22, 23]] > 1e-6)
bad_rows = bad_rows.unique()
all_rows = torch.arange(solar.size(0)).to(device)
good_rows = all_rows[~torch.isin(all_rows, bad_rows)]
# Drop the bad rows
solar = solar[good_rows]
# Remake the dataloader
dataset = dataloader.dataset[good_rows]
dataloader = torch.utils.data.DataLoader(dataset, batch_size=params['batch_size'], shuffle=True)
solar_forcing = SolarMultForcing()
solar_add_forcing = SolarAddForcing()
wind_forcing = WindMultForcing()
# mult_forcing = {"SOLAR": solar_forcing,
# "WIND": wind_forcing}
# add_forcing = {"SOLAR": solar_add_forcing}
initial_condition = SampledInitialCondition(data=dataloader.dataset[..., [1, 2]])
aux_solar_initial_condition = RandomInitialCondition(dim=1, low=0.2, high=0.8)
all_ic = ColumnwiseInitialCondition([
(initial_condition, [0, 1]),
(aux_solar_initial_condition, [2])
], device=device)
state_size = len(params['variables'])
G = Generator(initial_condition=all_ic,
state_size=state_size,
mlp_size=params['gen_mlp_size'],
num_layers=params['gen_num_layers'],
time_steps=params['time_series_length'],
varnames=params["variables"],
# mult_forcing=mult_forcing,
# add_forcing=add_forcing,
dawn_dusk_sampler=DawnDuskSampler(solar)).to(device)
D = DiscriminatorSimple(data_size=state_size,
time_size=params['time_series_length'],
num_layers=5,
num_units=200).to(device)
optimizer_G = Adam(G.parameters(), lr=params['gen_lr'], betas=params['gen_betas'])
optimizer_D = Adam(D.parameters(), lr=params['dis_lr'], betas=params['dis_betas'])
if warm_start and params_file is not None:
# load the models based on the model naming scheme for CNN models:
# saved_models/cnn/cnn_{gen/dis}_{ISO}_{var1}{var2}...{varn}.pt
# where {var1}...{varn} are the lowercase first letters of the variable names. This variable
# naming scheme isn't ideal since there can be collisions, but for the variables we're using
# it should be fine.
G.load_state_dict(torch.load(f'saved_models/sde/sde_gen_{params["ISO"]}_{"".join([v.lower()[0] for v in params["variables"]])}.pt'))
D.load_state_dict(torch.load(f'saved_models/sde/sde_dis_{params["ISO"]}_{"".join([v.lower()[0] for v in params["variables"]])}.pt'))
plotter = SDETrainingPlotter(['G', 'D'], varnames=params['variables'], transformer=transformer)
trainer = WGANGPTrainer(G, D, optimizer_G, optimizer_D,
plotter=plotter,
device=device,
penalty_weight=params['gp_weight'])
plot_every = max(1, params['epochs'] // 100)
print_every = max(1, params['epochs'] // 30)
trainer.train(data_loader=dataloader, epochs=params['epochs'], plot_every=plot_every, print_every=print_every)
if no_save:
return
# Save the trained models, parameters, and visualizations
dirname = 'saved_models/sde/'
if not os.path.exists(dirname):
os.makedirs(dirname)
# Save training visualizations
iso = params['ISO']
varnames_abbrev = ''.join([v.lower()[0] for v in params['variables']])
trainer.save_training_gif(dirname + f'training_sde_noembed_{iso}_{varnames_abbrev}.gif')
# Saving individual frames from the GIF. We need to be careful to not save a ton of frames.
save_every = len(plotter.frames) // 20 + 1 # will save at most 20 frames
plotter.save_frames(dirname + f'training_progress/training_sde_noembed_{iso}_{varnames_abbrev}.png',
save_every=save_every)
# Save models
torch.save(G.state_dict(), dirname + f'sde_gen_{iso}_{varnames_abbrev}.pt')
torch.save(D.state_dict(), dirname + f'sde_dis_{iso}_{varnames_abbrev}.pt')
# Save parameters
params['total_epochs_trained'] += params['epochs']
params['model_save_datetime'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
# reuse the params_file name if it was specified, otherwise use the default naming scheme
filename = params_file if params_file else dirname + f'params_sde_{iso}_{varnames_abbrev}.json'
with open(filename, 'w') as f:
json.dump(params, f)
if __name__ == '__main__':
fire.Fire(train_sdegan)