1140527 meeting

 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

wavenet:
  # WaveNet model parameters
  input_channels: 26  # Number of input channels
  output_channels: 26  # Number of output channels
  residual_layers: 36  # Number of residual layers
  residual_channels: 256  # Number of channels in residual blocks
  skip_channels: 256  # Number of channels in skip connections

  # Diffusion step embedding dimensions
  diffusion_step_embed_dim_input: 128  # Input dimension
  diffusion_step_embed_dim_hidden: 512  # Middle dimension
  diffusion_step_embed_dim_output: 512  # Output dimension

  # Structured State Spaces sequence model (S4) configurations
  s4_max_sequence_length: 2000  # Maximum sequence length
  s4_state_dim: 64  # State dimension
  s4_dropout: 0.0  # Dropout rate
  s4_bidirectional: true  # Whether to use bidirectional layers
  s4_use_layer_norm: true  # Whether to use layer normalization

diffusion:
  # Diffusion model parameters
  T: 200  # Number of diffusion steps
  beta_0: 0.0001  # Initial beta value
  beta_T: 0.02  # Final beta value

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

# Training configuration
batch_size: 1  # Batch size
output_directory: "./results/real_time/checkpoints/test"  # Output directory for checkpoints and logs
ckpt_iter: "max"  # Checkpoint mode (max or min)
iters_per_ckpt: 100  # Checkpoint frequency (number of epochs)
iters_per_logging: 100  # Log frequency (number of iterations)
n_iters: 20000  # Maximum number of iterations
learning_rate: 0.001  # Learning rate

# Additional training settings
only_generate_missing: true  # Generate missing values only
use_model: 2  # Model to use for training
masking: "forecast"  # Masking strategy for missing values
missing_k: 200  # Number of missing values

# Data paths
data:
  train_path: "./datasets/real_time/pollutants_train.npy"  # Path to training data
  location_path: "./datasets/real_time/locations.npy"  # Path to location data

  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

import argparse
import os
from typing import Optional, Union

import numpy as np
import torch
import yaml

from sssd.core.model_specs import MODEL_PATH_FORMAT, setup_model
from sssd.data.utils import get_dataloader
from sssd.training.trainer import DiffusionTrainer
from sssd.utils.logger import setup_logger
from sssd.utils.utils import calc_diffusion_hyperparams, display_current_time

LOGGER = setup_logger()


def fetch_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-m",
        "--model_config",
        type=str,
        default="configs/model.yaml",
        help="Model configuration",
    )
    parser.add_argument(
        "-t",
        "--training_config",
        type=str,
        default="configs/training.yaml",
        help="Training configuration",
    )
    return parser.parse_args()


def setup_output_directory(
    model_config: dict,
    training_config: dict,
) -> str:
    # Build output directory
    local_path = MODEL_PATH_FORMAT.format(
        T=model_config["diffusion"]["T"],
        beta_0=model_config["diffusion"]["beta_0"],
        beta_T=model_config["diffusion"]["beta_T"],
    )
    output_directory = os.path.join(training_config["output_directory"], local_path)

    if not os.path.isdir(output_directory):
        os.makedirs(output_directory)
        os.chmod(output_directory, 0o775)
    LOGGER.info("Output directory %s", output_directory)
    return output_directory

# New function to load locations
def get_locations_loader(location_path, device: torch.device) -> torch.Tensor:
    locations = np.load(location_path)
    locations_tensor = torch.from_numpy(locations).float().to(device)
    return locations_tensor

def run_job(
    model_config: dict,
    training_config: dict,
    device: Optional[Union[torch.device, str]],
) -> None:
    output_directory = setup_output_directory(model_config, training_config)
    dataloader = get_dataloader(
        training_config["data"]["train_path"],
        batch_size=training_config.get("batch_size"),
        device=device,
    )

    locations_loder = get_locations_loader(training_config["data"]["location_path"], device)

    diffusion_hyperparams = calc_diffusion_hyperparams(
        **model_config["diffusion"], device=device
    )
    net = setup_model(training_config["use_model"], model_config, device)

    LOGGER.info(display_current_time())
    trainer = DiffusionTrainer(
        dataloader=dataloader,
        diffusion_hyperparams=diffusion_hyperparams,
        net=net,
        device=device,
        output_directory=output_directory,
        ckpt_iter=training_config.get("ckpt_iter"),
        n_iters=training_config.get("n_iters"),
        iters_per_ckpt=training_config.get("iters_per_ckpt"),
        iters_per_logging=training_config.get("iters_per_logging"),
        learning_rate=training_config.get("learning_rate"),
        only_generate_missing=training_config.get("only_generate_missing"),
        masking=training_config.get("masking"),
        missing_k=training_config.get("missing_k"),
        batch_size=training_config.get("batch_size"),
        logger=LOGGER,
        locs=locations_loder,
    )
    trainer.train()

    LOGGER.info(display_current_time())


if __name__ == "__main__":
    args = fetch_args()

    with open(args.model_config, "rt") as f:
        model_config = yaml.safe_load(f.read())
    with open(args.training_config, "rt") as f:
        training_config = yaml.safe_load(f.read())

    LOGGER.info(f"Model spec: {model_config}")
    LOGGER.info(f"Training spec: {training_config}")

    if torch.cuda.device_count() > 0:
        LOGGER.info(f"Using {torch.cuda.device_count()} GPUs!")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    run_job(model_config, training_config, device)

  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
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254

import math

import torch
import torch.nn as nn

from sssd.core.layers.s4.s4_layer import S4Layer
from sssd.core.utils import calc_diffusion_step_embedding


def swish(x):
    return x * torch.sigmoid(x)


class Conv(nn.Module):
    def __init__(self, input_channels, output_channels, kernel_size=3, dilation=1):
        super().__init__()
        self.padding = dilation * (kernel_size - 1) // 2
        self.conv = nn.Conv1d(
            input_channels,
            output_channels,
            kernel_size,
            dilation=dilation,
            padding=self.padding,
        )
        self.conv = nn.utils.parametrizations.weight_norm(self.conv)
        nn.init.kaiming_normal_(self.conv.weight)

    def forward(self, x):
        out = self.conv(x)
        return out


class ZeroConv1d(nn.Module):
    def __init__(self, in_channel, out_channel):
        super(ZeroConv1d, self).__init__()
        self.conv = nn.Conv1d(in_channel, out_channel, kernel_size=1, padding=0)
        self.conv.weight.data.zero_()
        self.conv.bias.data.zero_()

    def forward(self, x):
        out = self.conv(x)
        return out


class ResidualBlock(nn.Module):
    def __init__(
        self,
        residual_channels,
        skip_channels,
        diffusion_step_embed_dim_output,
        input_channels,
        s4_max_sequence_length,
        s4_state_dim,
        s4_dropout,
        s4_bidirectional,
        s4_use_layer_norm,
    ):
        super().__init__()
        self.residual_channels = residual_channels

        self.fc_t = nn.Linear(diffusion_step_embed_dim_output, self.residual_channels)

        self.S41 = S4Layer(
            features=3 * self.residual_channels,
            lmax=s4_max_sequence_length,
            N=s4_state_dim,
            dropout=s4_dropout,
            bidirectional=s4_bidirectional,
            layer_norm=s4_use_layer_norm,
        )

        self.conv_layer = Conv(
            self.residual_channels, 3 * self.residual_channels, kernel_size=3
        )

        self.S42 = S4Layer(
            features=3 * self.residual_channels,
            lmax=s4_max_sequence_length,
            N=s4_state_dim,
            dropout=s4_dropout,
            bidirectional=s4_bidirectional,
            layer_norm=s4_use_layer_norm,
        )

        self.cond_conv = Conv(
            3 * input_channels, 3 * self.residual_channels, kernel_size=1
        )
        self.cond_bn = nn.BatchNorm1d(3 * self.residual_channels)
        self.res_conv = nn.Conv1d(residual_channels, residual_channels, kernel_size=1)
        self.res_conv = nn.utils.parametrizations.weight_norm(self.res_conv)
        nn.init.kaiming_normal_(self.res_conv.weight)

        self.skip_conv = nn.Conv1d(residual_channels, skip_channels, kernel_size=1)
        self.skip_conv = nn.utils.parametrizations.weight_norm(self.skip_conv)
        nn.init.kaiming_normal_(self.skip_conv.weight)

    def forward(self, input_data):
        x, cond, diffusion_step_embed = input_data
        h = x
        B, C, L = x.shape
        assert C == self.residual_channels

        part_t = self.fc_t(diffusion_step_embed)
        part_t = part_t.view([B, self.residual_channels, 1])
        h = h + part_t

        h = self.conv_layer(h)
        h = self.S41(h.permute(2, 0, 1)).permute(1, 2, 0)

        assert cond is not None
        cond = self.cond_conv(cond)
        cond = self.cond_bn(cond)
        h += cond

        h = self.S42(h.permute(2, 0, 1)).permute(1, 2, 0)

        #out = torch.tanh(h[:, : self.residual_channels, :]) * torch.sigmoid(
        #    h[:, self.residual_channels :, :]
        #)
        gate, filter, _ = torch.chunk(h, 3, dim=1)  # 每段 shape = (B, 256, L)
        out = torch.tanh(gate) * torch.sigmoid(filter)

        res = self.res_conv(out)
        assert x.shape == res.shape
        skip = self.skip_conv(out)

        return (x + res) * math.sqrt(0.5), skip  # normalize for training stability


class ResidualGroup(nn.Module):
    def __init__(
        self,
        residual_channels,
        skip_channels,
        residual_layers,
        diffusion_step_embed_dim_input,
        diffusion_step_embed_dim_hidden,
        diffusion_step_embed_dim_output,
        input_channels,
        s4_max_sequence_length,
        s4_state_dim,
        s4_dropout,
        s4_bidirectional,
        s4_use_layer_norm,
        device="cuda",
    ):
        super(ResidualGroup, self).__init__()
        self.residual_layers = residual_layers
        self.diffusion_step_embed_dim_input = diffusion_step_embed_dim_input

        self.fc_t1 = nn.Linear(
            diffusion_step_embed_dim_input, diffusion_step_embed_dim_hidden
        )
        self.fc_t2 = nn.Linear(
            diffusion_step_embed_dim_hidden, diffusion_step_embed_dim_output
        )

        self.residual_blocks = nn.ModuleList()
        for n in range(self.residual_layers):
            self.residual_blocks.append(
                ResidualBlock(
                    residual_channels,
                    skip_channels,
                    diffusion_step_embed_dim_output=diffusion_step_embed_dim_output,
                    input_channels=input_channels,
                    s4_max_sequence_length=s4_max_sequence_length,
                    s4_state_dim=s4_state_dim,
                    s4_dropout=s4_dropout,
                    s4_bidirectional=s4_bidirectional,
                    s4_use_layer_norm=s4_use_layer_norm,
                )
            )

        self.device = device

    def forward(self, input_data):
        noise, conditional, diffusion_steps = input_data

        diffusion_step_embed = calc_diffusion_step_embedding(
            diffusion_steps, self.diffusion_step_embed_dim_input, device=self.device
        )
        diffusion_step_embed = swish(self.fc_t1(diffusion_step_embed))
        diffusion_step_embed = swish(self.fc_t2(diffusion_step_embed))

        h = noise
        skip = 0
        for n in range(self.residual_layers):
            h, skip_n = self.residual_blocks[n]((h, conditional, diffusion_step_embed))
            skip += skip_n

        return skip * math.sqrt(1.0 / self.residual_layers)


class SSSDS4Imputer(nn.Module):
    def __init__(
        self,
        input_channels,
        residual_channels,
        skip_channels,
        output_channels,
        residual_layers,
        diffusion_step_embed_dim_input,
        diffusion_step_embed_dim_hidden,
        diffusion_step_embed_dim_output,
        s4_max_sequence_length,
        s4_state_dim,
        s4_dropout,
        s4_bidirectional,
        s4_use_layer_norm,
        device="cuda",
    ):
        super().__init__()

        self.init_conv = nn.Sequential(
            nn.Conv1d(input_channels, residual_channels, kernel_size=1),
            nn.ReLU(),
        )

        self.residual_layer = ResidualGroup(
            residual_channels=residual_channels,
            skip_channels=skip_channels,
            residual_layers=residual_layers,
            diffusion_step_embed_dim_input=diffusion_step_embed_dim_input,
            diffusion_step_embed_dim_hidden=diffusion_step_embed_dim_hidden,
            diffusion_step_embed_dim_output=diffusion_step_embed_dim_output,
            input_channels=input_channels,
            s4_max_sequence_length=s4_max_sequence_length,
            s4_state_dim=s4_state_dim,
            s4_dropout=s4_dropout,
            s4_bidirectional=s4_bidirectional,
            s4_use_layer_norm=s4_use_layer_norm,
            device=device,
        )

        self.final_conv = nn.Sequential(
            nn.Conv1d(skip_channels, skip_channels, kernel_size=1),
            nn.ReLU(),
            ZeroConv1d(skip_channels, output_channels),
        )

    def forward(self, input_data):
        noise, conditional, mask, diffusion_steps, frk_basis = input_data

        # Handle mask and concatenate it to the conditional input
        conditional = conditional * mask
        conditional = torch.cat([conditional, mask.float(), frk_basis], dim=1)

        # Forward pass through the network
        x = noise  # Ensure x is 3D (B, C, L)
        x = self.init_conv(x)
        x = self.residual_layer((x, conditional, diffusion_steps))
        y = self.final_conv(x)

        return y

  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
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219

import math

import torch
import torch.nn as nn

from sssd.core.utils import calc_diffusion_step_embedding


def swish(x):
    return x * torch.sigmoid(x)


class Conv(nn.Module):
    def __init__(self, input_channels, output_channels, kernel_size=3, dilation=1):
        super(Conv, self).__init__()
        self.padding = dilation * (kernel_size - 1) // 2
        self.conv = nn.Conv1d(
            input_channels,
            output_channels,
            kernel_size,
            dilation=dilation,
            padding=self.padding,
        )
        self.conv = nn.utils.parametrizations.weight_norm(self.conv)
        nn.init.kaiming_normal_(self.conv.weight)

    def forward(self, x):
        out = self.conv(x)
        return out


class ZeroConv1d(nn.Module):
    def __init__(self, in_channel, out_channel):
        super(ZeroConv1d, self).__init__()
        self.conv = nn.Conv1d(in_channel, out_channel, kernel_size=1, padding=0)
        self.conv.weight.data.zero_()
        self.conv.bias.data.zero_()

    def forward(self, x):
        out = self.conv(x)
        return out


class Residual_block(nn.Module):
    def __init__(
        self,
        residual_channels,
        skip_channels,
        dilation,
        diffusion_step_embed_dim_output,
        input_channels,
    ):
        super(Residual_block, self).__init__()

        self.residual_channels = residual_channels
        # the layer-specific fc for diffusion step embedding
        self.fc_t = nn.Linear(diffusion_step_embed_dim_output, self.residual_channels)

        # dilated conv layer
        self.dilated_conv_layer = Conv(
            self.residual_channels,
            3 * self.residual_channels,
            kernel_size=3,
            dilation=dilation,
        )

        # add mel spectrogram upsampler and conditioner conv1x1 layer  (In adapted to S4 output)
        self.cond_conv = Conv(
            3 * input_channels, 3 * self.residual_channels, kernel_size=1
        )  # 80 is mel bands

        # residual conv1x1 layer, connect to next residual layer
        self.res_conv = nn.Conv1d(residual_channels, residual_channels, kernel_size=1)
        self.res_conv = nn.utils.parametrizations.weight_norm(self.res_conv)
        nn.init.kaiming_normal_(self.res_conv.weight)

        # skip conv1x1 layer, add to all skip outputs through skip connections
        self.skip_conv = nn.Conv1d(residual_channels, skip_channels, kernel_size=1)
        self.skip_conv = nn.utils.parametrizations.weight_norm(self.skip_conv)
        nn.init.kaiming_normal_(self.skip_conv.weight)

    def forward(self, input_data):
        x, cond, diffusion_step_embed = input_data
        h = x
        B, C, L = x.shape
        assert C == self.residual_channels

        part_t = self.fc_t(diffusion_step_embed)
        part_t = part_t.view([B, self.residual_channels, 1])
        h = h + part_t

        h = self.dilated_conv_layer(h)
        # add (local) conditioner
        assert cond is not None

        cond = self.cond_conv(cond)
        h += cond

        out = torch.tanh(h[:, : self.residual_channels, :]) * torch.sigmoid(
            h[:, self.residual_channels :, :]
        )

        res = self.res_conv(out)
        assert x.shape == res.shape
        skip = self.skip_conv(out)

        return (x + res) * math.sqrt(0.5), skip


class Residual_group(nn.Module):
    def __init__(
        self,
        residual_channels,
        skip_channels,
        residual_layers,
        dilation_cycle,
        diffusion_step_embed_dim_input,
        diffusion_step_embed_dim_hidden,
        diffusion_step_embed_dim_output,
        input_channels,
        device="cuda",
    ):
        super(Residual_group, self).__init__()
        self.residual_layers = residual_layers
        self.diffusion_step_embed_dim_input = diffusion_step_embed_dim_input

        self.fc_t1 = nn.Linear(
            diffusion_step_embed_dim_input, diffusion_step_embed_dim_hidden
        )
        self.fc_t2 = nn.Linear(
            diffusion_step_embed_dim_hidden, diffusion_step_embed_dim_output
        )

        self.residual_blocks = nn.ModuleList()
        for n in range(self.residual_layers):
            self.residual_blocks.append(
                Residual_block(
                    residual_channels,
                    skip_channels,
                    dilation=2 ** (n % dilation_cycle),
                    diffusion_step_embed_dim_output=diffusion_step_embed_dim_output,
                    input_channels=input_channels,
                )
            )
        self.device = device

    def forward(self, input_data):
        noise, conditional, diffusion_steps = input_data

        diffusion_step_embed = calc_diffusion_step_embedding(
            diffusion_steps, self.diffusion_step_embed_dim_input, device=self.device
        )
        diffusion_step_embed = swish(self.fc_t1(diffusion_step_embed))
        diffusion_step_embed = swish(self.fc_t2(diffusion_step_embed))

        h = noise
        skip = 0
        for n in range(self.residual_layers):
            h, skip_n = self.residual_blocks[n](
                (noise, conditional, diffusion_step_embed)
            )
            skip += skip_n

        return skip * math.sqrt(
            1.0 / self.residual_layers
        )  # normalize for training stability


class DiffWaveImputer(nn.Module):
    def __init__(
        self,
        input_channels,
        residual_channels,
        skip_channels,
        output_channels,
        residual_layers,
        dilation_cycle,
        diffusion_step_embed_dim_input,
        diffusion_step_embed_dim_hidden,
        diffusion_step_embed_dim_output,
        device="cuda",
    ):
        super(DiffWaveImputer, self).__init__()

        self.init_conv = nn.Sequential(
            Conv(input_channels, residual_channels, kernel_size=1), nn.ReLU()
        )

        self.residual_layer = Residual_group(
            residual_channels=residual_channels,
            skip_channels=skip_channels,
            residual_layers=residual_layers,
            dilation_cycle=dilation_cycle,
            diffusion_step_embed_dim_input=diffusion_step_embed_dim_input,
            diffusion_step_embed_dim_hidden=diffusion_step_embed_dim_hidden,
            diffusion_step_embed_dim_output=diffusion_step_embed_dim_output,
            input_channels=input_channels,
            device=device,
        )

        self.final_conv = nn.Sequential(
            Conv(skip_channels, skip_channels, kernel_size=1),
            nn.ReLU(),
            ZeroConv1d(skip_channels, output_channels),
        )

    def forward(self, input_data):

        noise, conditional, mask, diffusion_steps, frk_basis = input_data

        conditional = conditional * mask
        conditional = torch.cat([conditional, mask.float(), frk_basis], dim=1)

        x = noise
        x = self.init_conv(x)
        x = self.residual_layer((x, conditional, diffusion_steps))
        y = self.final_conv(x)

        return y

  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
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333

import logging
import os
from typing import Any, Dict, Optional, Union

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

from sssd.core.model_specs import MASK_FN
from sssd.training.utils import training_loss
from sssd.utils.logger import setup_logger
from sssd.utils.utils import find_max_epoch

LOGGER = setup_logger()


class DiffusionTrainer:
    """
    Train Diffusion Models

    Args:
        dataloader (DataLoader): The training dataloader.
        diffusion_hyperparams (Dict[str, Any]): Hyperparameters for the diffusion process.
        net (nn.Module): The neural network model to be trained.
        device (torch.device): The device to be used for training.
        output_directory (str): Directory to save model checkpoints.
        ckpt_iter (Optional[int, str]): The checkpoint iteration to be loaded; 'max' selects the maximum iteration.
        n_iters (int): Number of iterations to train.
        iters_per_ckpt (int): Number of iterations to save checkpoint.
        iters_per_logging (int): Number of iterations to save training log and compute validation loss.
        learning_rate (float): Learning rate for training.
        only_generate_missing (int): Option to generate missing portions of the signal only.
        masking (str): Type of masking strategy: 'mnr' for Missing Not at Random, 'bm' for Blackout Missing, 'rm' for Random Missing.
        missing_k (int): K missing time steps for each feature across the sample length.
        batch_size (int): Size of each training batch.
        logger (Optional[logging.Logger]): Logger object for logging, defaults to None.
    """

    def __init__(
        self,
        dataloader: DataLoader,
        diffusion_hyperparams: Dict[str, Any],
        net: nn.Module,
        device: Optional[Union[torch.device, str]],
        output_directory: str,
        ckpt_iter: Union[str, int],
        n_iters: int,
        iters_per_ckpt: int,
        iters_per_logging: int,
        learning_rate: float,
        only_generate_missing: int,
        masking: str,
        missing_k: int,
        batch_size: int,
        logger: Optional[logging.Logger] = None,
        locs: Optional[torch.Tensor] = None
    ) -> None:
        self.dataloader = dataloader
        self.diffusion_hyperparams = diffusion_hyperparams
        self.net = nn.DataParallel(net).to(device)
        self.device = device
        self.output_directory = output_directory
        self.ckpt_iter = ckpt_iter
        self.n_iters = n_iters
        self.iters_per_ckpt = iters_per_ckpt
        self.iters_per_logging = iters_per_logging
        self.learning_rate = learning_rate
        self.only_generate_missing = only_generate_missing
        self.masking = masking
        self.missing_k = missing_k
        self.writer = SummaryWriter(f"{output_directory}/log")
        self.batch_size = batch_size
        self.optimizer = torch.optim.Adam(self.net.parameters(), lr=self.learning_rate)
        self.logger = logger or LOGGER


        self.locs = locs
        frk = FRKBasis(self.locs, device=self.device)
        self.frk_basis = frk()
        batch_size_shape, time_shape, loc_shape = next(iter(dataloader))[0].shape
        
        print(f"self.frk_basis is None, initializing...")
        frk_linear = nn.Linear(self.frk_basis.shape[1], time_shape).to(self.device)
        loc_embed = nn.Embedding(loc_shape, time_shape).to(self.device)

        self.frk_basis = frk_linear(self.frk_basis)  # [26, 29] → [26, 2000]
        loc_ids = torch.arange(loc_shape, device=self.device)
        self.frk_basis += loc_embed(loc_ids)  # 加入地點特徵
        self.frk_basis = self.frk_basis.unsqueeze(0).expand(batch_size_shape, -1, -1)  # [1, 26, 2000]
        self.frk_basis = self.frk_basis.detach()
        print(f"frk unsqueeze shape: {self.frk_basis.shape}")
        #print(f"frk original shape: {self.frk_basis.shape}")






        if self.masking not in MASK_FN:
            raise KeyError(f"Please enter a correct masking, but got {self.masking}")

    def _load_checkpoint(self) -> None:
        if self.ckpt_iter == "max":
            self.ckpt_iter = find_max_epoch(self.output_directory)
        if self.ckpt_iter >= 0:
            try:
                model_path = os.path.join(
                    self.output_directory, f"{self.ckpt_iter}.pkl"
                )
                checkpoint = torch.load(model_path, map_location="cpu")

                self.net.load_state_dict(checkpoint["model_state_dict"])
                if "optimizer_state_dict" in checkpoint:
                    self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

                self.logger.info(
                    f"Successfully loaded model at iteration {self.ckpt_iter}"
                )
            except Exception as e:
                self.ckpt_iter = -1
                self.logger.error(f"No valid checkpoint model found. Error: {e}")
        else:
            self.ckpt_iter = -1
            self.logger.info(
                "No valid checkpoint model found, start training from initialization."
            )

    def _save_model(self, n_iter: int) -> None:
        if n_iter > 0 and n_iter % self.iters_per_ckpt == 0:
            torch.save(
                {
                    "model_state_dict": self.net.state_dict(),
                    "optimizer_state_dict": self.optimizer.state_dict(),
                },
                os.path.join(self.output_directory, f"{n_iter}.pkl"),
            )

    def _update_mask(self, batch: torch.Tensor) -> torch.Tensor:
        transposed_mask = MASK_FN[self.masking](batch[0], self.missing_k)
        return (
            transposed_mask.permute(1, 0)
            .repeat(batch.size()[0], 1, 1)
            .to(self.device, dtype=torch.float32)
        )

    def _train_per_epoch(self) -> torch.Tensor:
        for (batch,) in tqdm(self.dataloader):
            batch = batch.to(self.device)
            mask = self._update_mask(batch)
            loss_mask = ~mask.bool()

            batch = batch.permute(0, 2, 1)

            #print(f"batch shape: {batch.shape}")
            #print(f"frk shape: {self.frk_basis.shape}")
            #print(f"mask shape: {mask.shape}")

            frk_basis = self.frk_basis
            
            assert batch.size() == mask.size() == loss_mask.size()

            self.optimizer.zero_grad()
            loss = training_loss(
                model=self.net,
                loss_function=nn.MSELoss(),
                training_data=(batch, batch, mask, loss_mask, frk_basis),
                diffusion_parameters=self.diffusion_hyperparams,
                generate_only_missing=self.only_generate_missing,
                device=self.device,
            )
            loss.backward()
            self.optimizer.step()

        return loss

    def train(self) -> None:
        self._load_checkpoint()

        n_iter_start = (
            self.ckpt_iter + 2 if self.ckpt_iter == -1 else self.ckpt_iter + 1
        )
        self.logger.info(f"Start the {n_iter_start} iteration")

        for n_iter in range(n_iter_start, self.n_iters + 1):
            loss = self._train_per_epoch()
            self.writer.add_scalar("Train/Loss", loss.item(), n_iter)
            if n_iter % self.iters_per_logging == 0:
                self.logger.info(f"Iteration: {n_iter} \tLoss: { loss.item()}")
            self._save_model(n_iter)


class FRKBasis(nn.Module):
    def __init__(self, locs: torch.Tensor, k: int=None, device: Optional[Union[torch.device, str]]=None, standardize: bool=True, unbiased_std: bool=False):
        """
        Multi-resolution spline basis function.
        
        Args:
            locs: [N, d] tensor of spatial locations.
            k: Optional parameter controlling resolution.
            device: PyTorch device to use.
            standardize: If True, perform standardization on locs (Recommended).
            unbiased_std: If True, use unbiased standard deviation (dividing by N-1) (Recommended); 
                          if False, use biased standard deviation (dividing by N).
        """
        super().__init__()
        try:
            if device is None:
                self.device = torch.device('cpu')
                LOGGER.warning(f'Parameter "device" was not set. Default value "cpu" is used.')
            else:
                self.device = torch.device(device)
                LOGGER.info(f'Successfully using device "{self.device}".')
        except (TypeError, ValueError) as e:
            self.device = torch.device('cpu')
            LOGGER.warning(f'Parameter "device" is not a valid device ({device}). Default value "cpu" is used. Error: {e}')

        self.register_buffer("locs", locs.to(self.device))
        self.N, self.d = locs.shape

        # number of basis
        try:
            if k is None:
                self.k = self.N + self.d + 1
                LOGGER.warning(f'Parameter "k" was not set. Default value {self.k} is used.')
            elif 0 < k <= (self.N + self.d + 1):
                self.k = k
            else:
                self.k = self.N + self.d + 1
                LOGGER.warning(f'Parameter "k" is out of valid range, it should be between 1 to {self.k}. Default value {self.k} is used.')
        except TypeError:
            self.k = self.N + self.d + 1
            LOGGER.warning(f'Parameter "k" is not an integer, the type you input is "{type(k).__name__}." Default value {self.k} is used.')

        self.standardize = standardize
        if type(self.standardize) is not bool:
            self.standardize = True
            LOGGER.warning(f'Parameter "standardize" should be a boolean, the type you input is "{type(standardize).__name__}". Default value True is used.')

        self.unbiased_std = unbiased_std
        if type(self.unbiased_std) is not bool:
            self.unbiased_std = False
            LOGGER.warning(f'Parameter "unbiased_std" should be a boolean, the type you input is "{type(unbiased_std).__name__}". Default value False is used.')

    def _tps_phi(self, r: torch.Tensor) -> torch.Tensor:
        """
        Apply TPS radial basis based on input distance r
        """
        eps = 1e-10  # avoid log(0)
        if self.d == 1:
            return (r ** 3) / 12
        elif self.d == 2:
            return (r ** 2) * torch.log(r + eps) / (8 * torch.pi)
        elif self.d == 3:
            return -r / 8
        else:
            raise NotImplementedError(f"TPS not implemented for d = {self.d}")
        
    def _standardize(self, x: torch.Tensor, unbiased_std: Optional[bool] = None) -> torch.Tensor:
        if unbiased_std is None:
            unbiased_std = self.unbiased_std
        mean = x.mean(dim=0, keepdim=True)  # keepdim=True to keep the same shape
        std = x.std(dim=0, unbiased=unbiased_std, keepdim=True)  # unbiased=False for population std
        std[std == 0] = 1.0
        x = (x - mean) / std
        return x

    def forward(self) -> torch.Tensor:

        # Construct X = [1, x1, x2, ..., xd]
        ones = torch.ones(self.N, 1, device=self.device)
        X = torch.cat([ones, self.locs], dim=1)  # [N, d+1]

        if self.k == 1:
            return ones
        elif self.k <= self.d:
            if self.standardize:
                X[:, 1:self.k] = self._standardize(X[:, 1:self.k])
            return X[:, :self.k]
        

        # TPS basis
        dists = torch.cdist(self.locs, self.locs)  # Euclidean distances
        Phi = self._tps_phi(dists)  # [N, N]

        # Q = I - X(XᵗX)⁻¹Xᵗ
        try:
            XtX_inv = torch.inverse(X.T @ X)  # [(d+1), (d+1)]
        except RuntimeError as e:
            LOGGER.warning(f'torch.inverse failed due to {e}, this implies "X.T @ X" didn\'t have inverse. Using torch.linalg.pinv instead.')
            XtX_inv = torch.linalg.pinv(X.T @ X)
        Q = torch.eye(self.N, device=self.device) - X @ XtX_inv @ X.T

        # Eigen-decomposition
        G = Q @ Phi @ Q
        eigenvalues, eigenvectors = torch.linalg.eigh(G)  # ascending order

        # Filter out near-zero eigenvalues and sort descending
        valid = eigenvalues > 1e-6
        eigenvalues[~valid] *= -1.0  # make them positive
        eigenvectors[:, ~valid] *= -1.0  # make them positive
        idx_desc = torch.argsort(eigenvalues, descending=True)
        eigenvalues = eigenvalues[idx_desc]
        eigenvectors = eigenvectors[:, idx_desc]

        # Construct basis functions
        Basis_high = (Phi.T - Phi @ X @ XtX_inv @ X.T).T @ eigenvectors @ torch.diag(1.0 / eigenvalues).to(self.device)
        
        # # the for loop
        # Basis_high_1 = torch.zeros(self.N, self.N, device=self.locs.device)
        # for i in range(self.N):
        #     phi_i = Phi[i, :]
        #     x_i = X[i, :]
        #     Phi_XXtXinv_x = (phi_i.T - Phi @ X @ XtX_inv @ x_i.T).T
        #     for j in range(self.N):
        #         lambda_j = 1 / eigenvalues[j]
        #         v_j = eigenvectors[:, j]
        #         Basis_high_1[i, j] = lambda_j * Phi_XXtXinv_x @ v_j
        # print(f'Basis_high:\n{Basis_high}, \n\nBasis_high_1:\n{Basis_high_1}')
        # print(torch.mean((Basis_high - Basis_high_1)**2))

        # Concatenate [1, x1, ..., xd] and B_high
        Basis_low = X  # [N, d+1]
        F = torch.cat([Basis_low, Basis_high], dim=1)

        # Standardize
        if self.standardize:
            F[:, 1:] = self._standardize(F[:, 1:])

        # Output k basis
        F = F[:, :self.k]
        return F  # [N, d+1 + k]