rename Batch to DataSetIter in enas_trainer
@@ -14,7 +14,7 @@ except:
from ..core.utils import _pseudo_tqdm as tqdm
from ..core.trainer import Trainer
-from ..core.batch import Batch
+from ..core.batch import DataSetIter
from ..core.callback import CallbackManager, CallbackException
from ..core.dataset import DataSet
from ..core.utils import _move_dict_value_to_device
@@ -124,7 +124,7 @@ class ENASTrainer(Trainer):
len(self.train_data) % self.batch_size != 0)) * self.n_epochs
with inner_tqdm(total=total_steps, postfix='loss:{0:<6.5f}', leave=False, dynamic_ncols=True) as pbar:
avg_loss = 0
- data_iterator = Batch(self.train_data, batch_size=self.batch_size, sampler=self.sampler, as_numpy=False,
+ data_iterator = DataSetIter(self.train_data, batch_size=self.batch_size, sampler=self.sampler, as_numpy=False,
prefetch=self.prefetch)
for epoch in range(1, self.n_epochs + 1):
pbar.set_description_str(desc="Epoch {}/{}".format(epoch, self.n_epochs))
@@ -209,7 +209,7 @@ class ENASTrainer(Trainer):
total_loss = 0
train_idx = 0
avg_loss = 0
- data_iterator = Batch(self.train_data, batch_size=self.batch_size, sampler=self.sampler, as_numpy=False,
+ data_iterator = DataSetIter(self.train_data, batch_size=self.batch_size, sampler=self.sampler, as_numpy=False,
prefetch=self.prefetch)
for batch_x, batch_y in data_iterator:
@@ -262,7 +262,7 @@ class ENASTrainer(Trainer):
if not isinstance(entropies, np.ndarray):
entropies = entropies.data.cpu().numpy()
- data_iterator = Batch(self.dev_data, batch_size=self.batch_size, sampler=self.sampler, as_numpy=False,
+ data_iterator = DataSetIter(self.dev_data, batch_size=self.batch_size, sampler=self.sampler, as_numpy=False,
prefetch=self.prefetch)
for inputs, targets in data_iterator:
