(
self,
*,
model,
diffusion,
data,
batch_size,
microbatch,
lr,
ema_rate,
log_interval,
save_interval,
resume_checkpoint,
use_fp16=False,
fp16_scale_growth=1e-3,
schedule_sampler=None,
weight_decay=0.0,
lr_anneal_steps=0,
)
| 25 | |
| 26 | class TrainLoop: |
| 27 | def __init__( |
| 28 | self, |
| 29 | *, |
| 30 | model, |
| 31 | diffusion, |
| 32 | data, |
| 33 | batch_size, |
| 34 | microbatch, |
| 35 | lr, |
| 36 | ema_rate, |
| 37 | log_interval, |
| 38 | save_interval, |
| 39 | resume_checkpoint, |
| 40 | use_fp16=False, |
| 41 | fp16_scale_growth=1e-3, |
| 42 | schedule_sampler=None, |
| 43 | weight_decay=0.0, |
| 44 | lr_anneal_steps=0, |
| 45 | ): |
| 46 | self.model = model |
| 47 | self.diffusion = diffusion |
| 48 | self.data = data |
| 49 | self.batch_size = batch_size |
| 50 | self.microbatch = microbatch if microbatch > 0 else batch_size |
| 51 | self.lr = lr |
| 52 | self.ema_rate = ( |
| 53 | [ema_rate] |
| 54 | if isinstance(ema_rate, float) |
| 55 | else [float(x) for x in ema_rate.split(",")] |
| 56 | ) |
| 57 | self.log_interval = log_interval |
| 58 | self.save_interval = save_interval |
| 59 | self.resume_checkpoint = resume_checkpoint |
| 60 | self.use_fp16 = use_fp16 |
| 61 | self.fp16_scale_growth = fp16_scale_growth |
| 62 | self.schedule_sampler = schedule_sampler or UniformSampler(diffusion) |
| 63 | self.weight_decay = weight_decay |
| 64 | self.lr_anneal_steps = lr_anneal_steps |
| 65 | |
| 66 | tblog_dir = os.path.join(logger.get_current().get_dir(), "tblog") |
| 67 | self.tb = SummaryWriter(tblog_dir) |
| 68 | |
| 69 | self.step = 0 |
| 70 | self.resume_step = 0 |
| 71 | self.global_batch = self.batch_size # * dist.get_world_size() |
| 72 | |
| 73 | self.sync_cuda = th.cuda.is_available() |
| 74 | |
| 75 | self._load_and_sync_parameters() |
| 76 | self.mp_trainer = MixedPrecisionTrainer( |
| 77 | model=self.model, |
| 78 | use_fp16=self.use_fp16, |
| 79 | fp16_scale_growth=fp16_scale_growth, |
| 80 | ) |
| 81 | |
| 82 | self.opt = AdamW( |
| 83 | self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay |
| 84 | ) |
nothing calls this directly
no test coverage detected