How to load the model i trained and denoise from a given step?

[ohaiyoskibidi]

i'm using the training code offered by the author:

from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer

model = Unet(
dim = 64,
dim_mults = (1, 2, 4, 8),
flash_attn = True
)

diffusion = GaussianDiffusion(
model,
image_size = 128,
timesteps = 1000, # number of steps
sampling_timesteps = 250 # number of sampling timesteps (using ddim for faster inference [see citation for ddim paper])
)

trainer = Trainer(
diffusion,
'ImageDataset',
train_batch_size = 32,
train_lr = 8e-5,
train_num_steps = 700000, # total training steps
gradient_accumulate_every = 2, # gradient accumulation steps
ema_decay = 0.995, # exponential moving average decay
amp = True, # turn on mixed precision
calculate_fid = True # whether to calculate fid during training
)

trainer.train()

Now I'm trying to verify the model I trained in another 'denoising.py' file, let's say the model's save path is "./results/model-best.pt"
Should I type the code above again except trainer.train() and add trainer.load("best")?
And do that means after this I can just type some code such as

for t in tqdm(reversed(range(0, t_max+1)), desc = 'sampling loop time step', total = t_max+1):

img, x_start = diffusion.p_sample(img, t)

to start denoising from the step t_max and image img i want?
Excuse me for my bad expression.

[balakrish181]

Now I'm trying to verify the model I trained in another 'denoising.py' file, let's say the model's save path is "./results/model-best.pt" Should I type the code above again except trainer.train() and add trainer.load("best")?

Yes, you can load the saved weights like this. I implemented the same way and also, raised a Pull-Request for the same #356

And do that means after this I can just type some code such as

for t in tqdm(reversed(range(0, t_max+1)), desc = 'sampling loop time step', total = t_max+1):

img, x_start = diffusion.p_sample(img, t)

to start denoising from the step t_max and image img i want? Excuse me for my bad expression.

Yes, this will work. However, there is a function sample in diffusion class, which does this under the hood.

    @torch.inference_mode()
    def sample(self, batch_size = 16, return_all_timesteps = False):
        (h, w), channels = self.image_size, self.channels
        sample_fn = self.p_sample_loop if not self.is_ddim_sampling else self.ddim_sample
        return sample_fn((batch_size, channels, h, w), return_all_timesteps = return_all_timesteps)