FilesExpand file tree

main

/

test_training.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

86 lines (73 loc) · 3.32 KB

main

/

test_training.py

Copy path

Top

File metadata and controls

• Code
• Blame

86 lines (73 loc) · 3.32 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

# coding=utf-8

# Copyright 2022 HuggingFace Inc.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import unittest

import torch

from diffusers import DDIMScheduler, DDPMScheduler, UNet2DModel

from diffusers.testing_utils import slow

from diffusers.training_utils import set_seed

torch.backends.cuda.matmul.allow_tf32 = False

class TrainingTests(unittest.TestCase):

def get_model_optimizer(self, resolution=32):

set_seed(0)

model = UNet2DModel(sample_size=resolution, in_channels=3, out_channels=3)

optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)

return model, optimizer

@slow

def test_training_step_equality(self):

device = "cpu" # ensure full determinism without setting the CUBLAS_WORKSPACE_CONFIG env variable

ddpm_scheduler = DDPMScheduler(

num_train_timesteps=1000,

beta_start=0.0001,

beta_end=0.02,

beta_schedule="linear",

clip_sample=True,

)

ddim_scheduler = DDIMScheduler(

num_train_timesteps=1000,

beta_start=0.0001,

beta_end=0.02,

beta_schedule="linear",

clip_sample=True,

)

assert ddpm_scheduler.config.num_train_timesteps == ddim_scheduler.config.num_train_timesteps

# shared batches for DDPM and DDIM

set_seed(0)

clean_images = [torch.randn((4, 3, 32, 32)).clip(-1, 1).to(device) for _ in range(4)]

noise = [torch.randn((4, 3, 32, 32)).to(device) for _ in range(4)]

timesteps = [torch.randint(0, 1000, (4,)).long().to(device) for _ in range(4)]

# train with a DDPM scheduler

model, optimizer = self.get_model_optimizer(resolution=32)

model.train().to(device)

for i in range(4):

optimizer.zero_grad()

ddpm_noisy_images = ddpm_scheduler.add_noise(clean_images[i], noise[i], timesteps[i])

ddpm_noise_pred = model(ddpm_noisy_images, timesteps[i]).sample

loss = torch.nn.functional.mse_loss(ddpm_noise_pred, noise[i])

loss.backward()

optimizer.step()

del model, optimizer

# recreate the model and optimizer, and retry with DDIM

model, optimizer = self.get_model_optimizer(resolution=32)

model.train().to(device)

for i in range(4):

optimizer.zero_grad()

ddim_noisy_images = ddim_scheduler.add_noise(clean_images[i], noise[i], timesteps[i])

ddim_noise_pred = model(ddim_noisy_images, timesteps[i]).sample

loss = torch.nn.functional.mse_loss(ddim_noise_pred, noise[i])

loss.backward()

optimizer.step()

del model, optimizer

self.assertTrue(torch.allclose(ddpm_noisy_images, ddim_noisy_images, atol=1e-5))

self.assertTrue(torch.allclose(ddpm_noise_pred, ddim_noise_pred, atol=1e-5))