20250930 meeting

Published 2025-09-29 Updated 2025-10-05 434 words 3 minutes views

Contents

In this time, we want to know how to embed the autoFRK model into SSSD model. Now, we decided to make it by:

The code planned to be added for autoFRK is as follows, located in the /sssd/training/trainer.py file of the SSSD model.

 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
class DiffusionTrainer:
    ...
    
    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)
            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),
                diffusion_parameters=self.diffusion_hyperparams,
                generate_only_missing=self.only_generate_missing,
                device=self.device,
            )
            loss.backward()
            self.optimizer.step()
        # autoFRK step
        # calculate loss from autoFRK
        # backward step
            #loss.backward()
            #self.optimizer.step()
        return loss

To integrate the autoFRK model, the parts to be checked here are as follows:

Verify that the loss return value is either a scalar or a matrix (initially, it should be a single scalar).
Check how the current parameters are used for SSSD predictions, to facilitate subsequent interpolation by autoFRK after prediction.
Embed the autoFRK module and return interpolations for unknown locations while simultaneously computing the loss.
Confirm whether loss.backward() and self.optimizer.step() steps are necessary (they likely are, otherwise parameters cannot be updated via gradients).

The current modification idea will first be handled by ChatGPT, with the modification outline written as follows:

 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
def _train_per_epoch(self) -> torch.Tensor:
    # Step 1: Standard batch training
    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)
        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),
            diffusion_parameters=self.diffusion_hyperparams,
            generate_only_missing=self.only_generate_missing,
            device=self.device,
        )
        loss.backward()
        self.optimizer.step()

    # Step 2: Perform full sampling prediction using updated parameters
    with torch.no_grad():
        generated_series = sampling(
            net=self.net,
            size=batch.shape,  # Use the shape of the last batch
            diffusion_hyperparams=self.diffusion_hyperparams,
            cond=batch,
            mask=mask,
            only_generate_missing=self.only_generate_missing,
            device=self.device,
        )

    # Step 3: Fill unknown regions using autoFRK
    # TODO: Pass generated_series, batch, mask to autoFRK
    # imputed_series = autoFRK(generated_series, mask, ...)

    imputed_series = generated_series  # Temporarily use the sampling result directly

    # Step 4: Compute loss for the unknown regions using the imputed result
    self.optimizer.zero_grad()
    loss = nn.MSELoss()(
        imputed_series[~mask.bool()],
        batch[~mask.bool()]
    )

    # Step 5: Update parameters
    loss.backward()
    self.optimizer.step()

    return loss

Contents

20250930 meeting

Related Posts