Coverage for widitapp/diffusion/timestep_sampler.py: 100.00%

1# Modified from OpenAI's diffusion repos 

2# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py 

3# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion 

4# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py 

5 

6from abc import ABC, abstractmethod 

7 

8import numpy as np 

9import torch as th 

10import torch.distributed as dist 

11 

12 

13def create_named_schedule_sampler(name, diffusion): 

14 """ 

15 Create a ScheduleSampler from a library of pre-defined samplers. 

16 :param name: the name of the sampler. 

17 :param diffusion: the diffusion object to sample for. 

18 """ 

19 if name == "uniform": 

20 return UniformSampler(diffusion) 

21 elif name == "loss-second-moment": 

22 return LossSecondMomentResampler(diffusion) 

23 else: 

24 raise NotImplementedError(f"unknown schedule sampler: {name}") 

25 

26 

27class ScheduleSampler(ABC): 

28 """ 

29 A distribution over timesteps in the diffusion process, intended to reduce 

30 variance of the objective. 

31 By default, samplers perform unbiased importance sampling, in which the 

32 objective's mean is unchanged. 

33 However, subclasses may override sample() to change how the resampled 

34 terms are reweighted, allowing for actual changes in the objective. 

35 """ 

36 

37 @abstractmethod 

38 def weights(self): 

39 """ 

40 Get a numpy array of weights, one per diffusion step. 

41 The weights needn't be normalized, but must be positive. 

42 """ 

43 

44 def sample(self, batch_size, device): 

45 """ 

46 Importance-sample timesteps for a batch. 

47 :param batch_size: the number of timesteps. 

48 :param device: the torch device to save to. 

49 :return: a tuple (timesteps, weights): 

50 - timesteps: a tensor of timestep indices. 

51 - weights: a tensor of weights to scale the resulting losses. 

52 """ 

53 w = self.weights() 

54 p = w / np.sum(w) 

55 indices_np = np.random.choice(len(p), size=(batch_size,), p=p) 

56 indices = th.from_numpy(indices_np).long().to(device) 

57 weights_np = 1 / (len(p) * p[indices_np]) 

58 weights = th.from_numpy(weights_np).float().to(device) 

59 return indices, weights 

60 

61 

62class UniformSampler(ScheduleSampler): 

63 def __init__(self, diffusion): 

64 self.diffusion = diffusion 

65 self._weights = np.ones([diffusion.num_timesteps]) 

66 

67 def weights(self): 

68 return self._weights 

69 

70 

71class LossAwareSampler(ScheduleSampler): 

72 def update_with_local_losses(self, local_ts, local_losses): 

73 """ 

74 Update the reweighting using losses from a model. 

75 Call this method from each rank with a batch of timesteps and the 

76 corresponding losses for each of those timesteps. 

77 This method will perform synchronization to make sure all of the ranks 

78 maintain the exact same reweighting. 

79 :param local_ts: an integer Tensor of timesteps. 

80 :param local_losses: a 1D Tensor of losses. 

81 """ 

82 batch_sizes = [ 

83 th.tensor([0], dtype=th.int32, device=local_ts.device) 

84 for _ in range(dist.get_world_size()) 

85 ] 

86 dist.all_gather( 

87 batch_sizes, 

88 th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), 

89 ) 

90 

91 # Pad all_gather batches to be the maximum batch size. 

92 batch_sizes = [x.item() for x in batch_sizes] 

93 max_bs = max(batch_sizes) 

94 

95 timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] 

96 loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] 

97 dist.all_gather(timestep_batches, local_ts) 

98 dist.all_gather(loss_batches, local_losses) 

99 timesteps = [ 

100 x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] 

101 ] 

102 losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] 

103 self.update_with_all_losses(timesteps, losses) 

104 

105 @abstractmethod 

106 def update_with_all_losses(self, ts, losses): 

107 """ 

108 Update the reweighting using losses from a model. 

109 Sub-classes should override this method to update the reweighting 

110 using losses from the model. 

111 This method directly updates the reweighting without synchronizing 

112 between workers. It is called by update_with_local_losses from all 

113 ranks with identical arguments. Thus, it should have deterministic 

114 behavior to maintain state across workers. 

115 :param ts: a list of int timesteps. 

116 :param losses: a list of float losses, one per timestep. 

117 """ 

118 

119 

120class LossSecondMomentResampler(LossAwareSampler): 

121 def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): 

122 self.diffusion = diffusion 

123 self.history_per_term = history_per_term 

124 self.uniform_prob = uniform_prob 

125 self._loss_history = np.zeros( 

126 [diffusion.num_timesteps, history_per_term], dtype=np.float64 

127 ) 

128 self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=int) 

129 

130 def weights(self): 

131 if not self._warmed_up(): 

132 return np.ones([self.diffusion.num_timesteps], dtype=np.float64) 

133 weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1)) 

134 weights /= np.sum(weights) 

135 weights *= 1 - self.uniform_prob 

136 weights += self.uniform_prob / len(weights) 

137 return weights 

138 

139 def update_with_all_losses(self, ts, losses): 

140 for t, loss in zip(ts, losses): 

141 if self._loss_counts[t] == self.history_per_term: 

142 # Shift out the oldest loss term. 

143 self._loss_history[t, :-1] = self._loss_history[t, 1:] 

144 self._loss_history[t, -1] = loss 

145 else: 

146 self._loss_history[t, self._loss_counts[t]] = loss 

147 self._loss_counts[t] += 1 

148 

149 def _warmed_up(self): 

150 return (self._loss_counts == self.history_per_term).all()