Drops labels to enable classifier-free guidance.
(self, labels, force_drop_ids=None)
| 99 | self.dropout_prob = dropout_prob |
| 100 | |
| 101 | def token_drop(self, labels, force_drop_ids=None): |
| 102 | """ |
| 103 | Drops labels to enable classifier-free guidance. |
| 104 | """ |
| 105 | if force_drop_ids is None: |
| 106 | drop_ids = torch.rand( |
| 107 | labels.shape[0], device=labels.device |
| 108 | ) < self.dropout_prob |
| 109 | drop_ids = drop_ids.cuda() |
| 110 | dist.broadcast( |
| 111 | drop_ids, |
| 112 | fs_init.get_model_parallel_src_rank(), |
| 113 | fs_init.get_model_parallel_group(), |
| 114 | ) |
| 115 | drop_ids = drop_ids.to(labels.device) |
| 116 | else: |
| 117 | drop_ids = force_drop_ids == 1 |
| 118 | labels = torch.where(drop_ids, self.num_classes, labels) |
| 119 | return labels |
| 120 | |
| 121 | def forward(self, labels, train, force_drop_ids=None): |
| 122 | use_dropout = self.dropout_prob > 0 |