mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Try to make diffusion fid more deterministic

Browse Source
This commit is contained in:
James Betker 2021-06-14 09:27:43 -06:00
parent 5b4f86293f
commit 6b32c87dcb
3 changed files with 25 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
codes/data/image_corruptor.py
View File

@ -8,12 +8,6 @@ from PIL import Image
from io import BytesIO
# Feeds a random uniform through a cosine distribution to slightly bias corruptions towards "uncorrupted".
# Return is on [0,1] with a bias towards 0.
def get_rand():
r = random.random()
return 1 - cos(r * pi / 2)
# Get a rough visualization of the above distribution. (Y-axis is meaningless, just spreads data)
'''
if __name__ == '__main__':
@ -28,12 +22,26 @@ if __name__ == '__main__':
# options.
class ImageCorruptor:
def __init__(self, opt):
self.opt = opt
self.blur_scale = opt['corruption_blur_scale'] if 'corruption_blur_scale' in opt.keys() else 1
self.fixed_corruptions = opt['fixed_corruptions'] if 'fixed_corruptions' in opt.keys() else []
self.num_corrupts = opt['num_corrupts_per_image'] if 'num_corrupts_per_image' in opt.keys() else 0
if self.num_corrupts == 0:
return
self.random_corruptions = opt['random_corruptions'] if 'random_corruptions' in opt.keys() else []
self.reset_random()
def reset_random(self):
if 'random_seed' in self.opt.keys():
self.rand = random.Random(self.opt['random_seed'])
else:
self.rand = random.Random()
# Feeds a random uniform through a cosine distribution to slightly bias corruptions towards "uncorrupted".
# Return is on [0,1] with a bias towards 0.
def get_rand(self):
r = self.rand.random()
return 1 - cos(r * pi / 2)
def corrupt_images(self, imgs, return_entropy=False):
if self.num_corrupts == 0 and not self.fixed_corruptions:
@ -53,10 +61,10 @@ class ImageCorruptor:
applied_augs = augmentations + self.fixed_corruptions
for img in imgs:
for aug in augmentations:
r = get_rand()
r = self.get_rand()
img = self.apply_corruption(img, aug, r, applied_augs)
for aug in self.fixed_corruptions:
r = get_rand()
r = self.get_rand()
img = self.apply_corruption(img, aug, r, applied_augs)
entropy.append(r)
corrupted_imgs.append(img)

3
codes/data/image_folder_dataset.py
View File

@ -124,6 +124,9 @@ class ImageFolderDataset:
ls, ent = self.corruptor.corrupt_images(ls, return_entropy=True)
return ls, ent
def reset_random(self):
self.corruptor.reset_random()
def __len__(self):
return self.len

8
codes/trainer/eval/sr_diffusion_fid.py
View File

@ -15,6 +15,7 @@ from trainer.injectors.gaussian_diffusion_injector import GaussianDiffusionInfer
from utils.util import opt_get
# Performs a FID evaluation on a diffusion network
class SrDiffusionFidEvaluator(evaluator.Evaluator):
def __init__(self, model, opt_eval, env):
super().__init__(model, opt_eval, env)
@ -24,15 +25,18 @@ class SrDiffusionFidEvaluator(evaluator.Evaluator):
self.dataset = create_dataset(opt_eval['dataset'])
self.fid_real_samples = opt_eval['dataset']['paths'] # This is assumed to exist for the given dataset.
assert isinstance(self.fid_real_samples, str)
self.dataloader = DataLoader(self.dataset, self.batch_sz, shuffle=False, num_workers=1)
self.gd = GaussianDiffusionInferenceInjector(opt_eval['diffusion_params'], env)
self.out_key = opt_eval['diffusion_params']['out']
def perform_eval(self):
# Attempt to make the dataset deterministic.
self.dataset.reset_random()
dataloader = DataLoader(self.dataset, self.batch_sz, shuffle=False, num_workers=0)
fid_fake_path = osp.join(self.env['base_path'], "..", "fid", str(self.env["step"]))
os.makedirs(fid_fake_path, exist_ok=True)
counter = 0
for batch in tqdm(self.dataloader):
for batch in tqdm(dataloader):
batch = {k: v.to(self.env['device']) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
gen = self.gd(batch)[self.out_key]