mirror of
https://github.com/wassname/denoising-diffusion-pytorch.git
synced 2026-09-10 12:01:08 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
edc6d5136a | ||
|
|
207d23af86 | ||
|
|
cf4f44909f | ||
|
|
a7f2d670bb | ||
|
|
6ab29d5cea | ||
|
|
ddc31bc489 | ||
|
|
e872ec3618 |
@@ -42,7 +42,7 @@ diffusion = GaussianDiffusion(
|
|||||||
loss_type = 'l1' # L1 or L2
|
loss_type = 'l1' # L1 or L2
|
||||||
)
|
)
|
||||||
|
|
||||||
training_images = torch.randn(8, 3, 128, 128) # images are normalized from 0 to 1
|
training_images = torch.rand(8, 3, 128, 128) # images are normalized from 0 to 1
|
||||||
loss = diffusion(training_images)
|
loss = diffusion(training_images)
|
||||||
loss.backward()
|
loss.backward()
|
||||||
# after a lot of training
|
# after a lot of training
|
||||||
@@ -124,7 +124,7 @@ diffusion = GaussianDiffusion1D(
|
|||||||
objective = 'pred_v'
|
objective = 'pred_v'
|
||||||
)
|
)
|
||||||
|
|
||||||
training_seq = torch.randn(8, 32, 128) # features are normalized from 0 to 1
|
training_seq = torch.rand(8, 32, 128) # features are normalized from 0 to 1
|
||||||
loss = diffusion(training_seq)
|
loss = diffusion(training_seq)
|
||||||
loss.backward()
|
loss.backward()
|
||||||
|
|
||||||
@@ -248,3 +248,13 @@ sampled_seq.shape # (4, 32, 128)
|
|||||||
volume = {abs/2207.12598}
|
volume = {abs/2207.12598}
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
|
```bibtex
|
||||||
|
@article{Sunkara2022NoMS,
|
||||||
|
title = {No More Strided Convolutions or Pooling: A New CNN Building Block for Low-Resolution Images and Small Objects},
|
||||||
|
author = {Raja Sunkara and Tie Luo},
|
||||||
|
journal = {ArXiv},
|
||||||
|
year = {2022},
|
||||||
|
volume = {abs/2208.03641}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|||||||
@@ -23,6 +23,8 @@ from ema_pytorch import EMA
|
|||||||
|
|
||||||
from accelerate import Accelerator
|
from accelerate import Accelerator
|
||||||
|
|
||||||
|
from denoising_diffusion_pytorch.version import __version__
|
||||||
|
|
||||||
# constants
|
# constants
|
||||||
|
|
||||||
ModelPrediction = namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
|
ModelPrediction = namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
|
||||||
@@ -86,7 +88,10 @@ def Upsample(dim, dim_out = None):
|
|||||||
)
|
)
|
||||||
|
|
||||||
def Downsample(dim, dim_out = None):
|
def Downsample(dim, dim_out = None):
|
||||||
return nn.Conv2d(dim, default(dim_out, dim), 4, 2, 1)
|
return nn.Sequential(
|
||||||
|
Rearrange('b c (h p1) (w p2) -> b (c p1 p2) h w', p1 = 2, p2 = 2),
|
||||||
|
nn.Conv2d(dim * 4, default(dim_out, dim), 1)
|
||||||
|
)
|
||||||
|
|
||||||
class WeightStandardizedConv2d(nn.Conv2d):
|
class WeightStandardizedConv2d(nn.Conv2d):
|
||||||
"""
|
"""
|
||||||
@@ -426,7 +431,7 @@ class GaussianDiffusion(nn.Module):
|
|||||||
beta_schedule = 'cosine',
|
beta_schedule = 'cosine',
|
||||||
p2_loss_weight_gamma = 0., # p2 loss weight, from https://arxiv.org/abs/2204.00227 - 0 is equivalent to weight of 1 across time - 1. is recommended
|
p2_loss_weight_gamma = 0., # p2 loss weight, from https://arxiv.org/abs/2204.00227 - 0 is equivalent to weight of 1 across time - 1. is recommended
|
||||||
p2_loss_weight_k = 1,
|
p2_loss_weight_k = 1,
|
||||||
ddim_sampling_eta = 1.
|
ddim_sampling_eta = 0.
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
assert not (type(self) == GaussianDiffusion and model.channels != model.out_dim)
|
assert not (type(self) == GaussianDiffusion and model.channels != model.out_dim)
|
||||||
@@ -805,8 +810,8 @@ class Trainer(object):
|
|||||||
if self.accelerator.is_main_process:
|
if self.accelerator.is_main_process:
|
||||||
self.ema = EMA(diffusion_model, beta = ema_decay, update_every = ema_update_every)
|
self.ema = EMA(diffusion_model, beta = ema_decay, update_every = ema_update_every)
|
||||||
|
|
||||||
self.results_folder = Path(results_folder)
|
self.results_folder = Path(results_folder)
|
||||||
self.results_folder.mkdir(exist_ok = True)
|
self.results_folder.mkdir(exist_ok = True)
|
||||||
|
|
||||||
# step counter state
|
# step counter state
|
||||||
|
|
||||||
@@ -825,7 +830,8 @@ class Trainer(object):
|
|||||||
'model': self.accelerator.get_state_dict(self.model),
|
'model': self.accelerator.get_state_dict(self.model),
|
||||||
'opt': self.opt.state_dict(),
|
'opt': self.opt.state_dict(),
|
||||||
'ema': self.ema.state_dict(),
|
'ema': self.ema.state_dict(),
|
||||||
'scaler': self.accelerator.scaler.state_dict() if exists(self.accelerator.scaler) else None
|
'scaler': self.accelerator.scaler.state_dict() if exists(self.accelerator.scaler) else None,
|
||||||
|
'version': __version__
|
||||||
}
|
}
|
||||||
|
|
||||||
torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))
|
torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))
|
||||||
@@ -843,6 +849,9 @@ class Trainer(object):
|
|||||||
self.opt.load_state_dict(data['opt'])
|
self.opt.load_state_dict(data['opt'])
|
||||||
self.ema.load_state_dict(data['ema'])
|
self.ema.load_state_dict(data['ema'])
|
||||||
|
|
||||||
|
if 'version' in data:
|
||||||
|
print(f"loading from version {data['version']}")
|
||||||
|
|
||||||
if exists(self.accelerator.scaler) and exists(data['scaler']):
|
if exists(self.accelerator.scaler) and exists(data['scaler']):
|
||||||
self.accelerator.scaler.load_state_dict(data['scaler'])
|
self.accelerator.scaler.load_state_dict(data['scaler'])
|
||||||
|
|
||||||
|
|||||||
@@ -414,7 +414,7 @@ class GaussianDiffusion1D(nn.Module):
|
|||||||
beta_schedule = 'cosine',
|
beta_schedule = 'cosine',
|
||||||
p2_loss_weight_gamma = 0.,
|
p2_loss_weight_gamma = 0.,
|
||||||
p2_loss_weight_k = 1,
|
p2_loss_weight_k = 1,
|
||||||
ddim_sampling_eta = 1.
|
ddim_sampling_eta = 0.
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.model = model
|
self.model = model
|
||||||
|
|||||||
@@ -0,0 +1 @@
|
|||||||
|
__version__ = '0.1.5'
|
||||||
@@ -1,9 +1,11 @@
|
|||||||
from setuptools import setup, find_packages
|
from setuptools import setup, find_packages
|
||||||
|
|
||||||
|
exec(open('denoising_diffusion_pytorch/version.py').read())
|
||||||
|
|
||||||
setup(
|
setup(
|
||||||
name = 'denoising-diffusion-pytorch',
|
name = 'denoising-diffusion-pytorch',
|
||||||
packages = find_packages(),
|
packages = find_packages(),
|
||||||
version = '0.31.1',
|
version = __version__,
|
||||||
license='MIT',
|
license='MIT',
|
||||||
description = 'Denoising Diffusion Probabilistic Models - Pytorch',
|
description = 'Denoising Diffusion Probabilistic Models - Pytorch',
|
||||||
author = 'Phil Wang',
|
author = 'Phil Wang',
|
||||||
|
|||||||
Reference in New Issue
Block a user