Generate information to be passed to ES callback. Args: idx_in_batch: Index in batch for per-sample mode, None indicates batch mode where we pass only the iteration number. lead: how many batches ahead is this job wrt actual iteration
(self, idx_in_batch, epoch_idx, batch_size=0, lead=0)
| 203 | ] |
| 204 | |
| 205 | def callback_args(self, idx_in_batch, epoch_idx, batch_size=0, lead=0): |
| 206 | """Generate information to be passed to ES callback. |
| 207 | |
| 208 | Args: |
| 209 | idx_in_batch: Index in batch for per-sample mode, None indicates batch mode where we |
| 210 | pass only the iteration number. |
| 211 | lead: how many batches ahead is this job wrt actual iteration |
| 212 | """ |
| 213 | if not self.accepts_arg: |
| 214 | return () |
| 215 | if idx_in_batch is not None: |
| 216 | arg = nvidia.dali.types.SampleInfo( |
| 217 | self.current_sample + idx_in_batch + batch_size * lead, |
| 218 | idx_in_batch, |
| 219 | self.current_iter + lead, |
| 220 | epoch_idx, |
| 221 | ) |
| 222 | elif self.batch_info: |
| 223 | arg = nvidia.dali.types.BatchInfo(self.current_iter + lead, epoch_idx) |
| 224 | else: |
| 225 | arg = self.current_iter + lead |
| 226 | return (arg,) |
| 227 | |
| 228 | def reset_indices(self): |
| 229 | self.current_iter = 0 |
no outgoing calls
no test coverage detected