def __init__(
self,
config,
world_rank,
local_rank,
is_distributed,
use_gpu,
device,
use_fp16=False,
use_tqdm=False,
wrap_ddp=False,
add_dist_sampler=False,
scheduler_step_freq=None,
):
# You are not expected to override this method.
self._world_rank = world_rank
self._local_rank = local_rank
self._config = config
self._is_distributed = is_distributed
self._use_fp16 = choose_amp_backend(use_fp16, amp, apex_amp)
self._device = device
self._use_gpu = use_gpu and torch.cuda.is_available()
if tqdm is None and use_tqdm:
raise ValueError("tqdm must be installed to use tqdm in training.")
self._use_tqdm = use_tqdm
self.global_step = 0
self._wrap_ddp = wrap_ddp
self._add_dist_sampler = add_dist_sampler
self._scheduler_step_freq = scheduler_step_freq
self.timers = TimerCollection()
self.setup(config)
def __init__(self,
training_operator_cls,
config=None,
use_gpu=False,
serialize_data_creation=True,
use_fp16=False,
use_tqdm=False,
scheduler_step_freq=None):
self.training_operator_cls = training_operator_cls
self.config = {} if config is None else config
self.timers = utils.TimerCollection()
self.epochs = 0
self.training_operator = None
self.serialize_data_creation = serialize_data_creation
self.use_gpu = use_gpu
self.use_fp16 = choose_amp_backend(use_fp16, amp, apex_amp)
self.use_tqdm = use_tqdm
self.scheduler_step_freq = scheduler_step_freq
# Training and Validation iterators
self.train_iterator = None
self._should_reset_train_loader = True
self.val_iterator = None
self._should_reset_val_loader = True
