def collect_data(self):
if self.config.with_fit:
self.train_iter = ImageFolderDataIter(
root=self.config.image_dir,
data_shape=(3, 128, 128),
label_shape=(2, ),
data_names=['data'],
label_names=['softmax_label'],
flag=1,
transform=lambda data, label:
(data.astype(np.float32) / 255, label),
batch_size=self.config.batch_size)
# data_iter = iter(self.train_iter)
# batch = next(data_iter)
# logging.info(f'[] label: {batch.label}')
# # sys.exit(0)
# self.train_dataset = ImageFolderDataset(
# root=config.image_dir,
# flag=1,
# transform=lambda data, label: (data.astype(np.float32)/255, label)
# )
# self.train_dataloader = mx.gluon.data.DataLoader(
# dataset=self.train_dataset,
# batch_size=self.config.batch_size,
# shuffle=True
# )
# self.eval_dataset = ImageFolderDataset(
# root=config.image_dir,
# flag=1,
# transform=lambda data, label: (data.astype(np.float32)/255, label)
# )
# self.eval_dataloader = mx.gluon.data.DataLoader(
# dataset=self.eval_dataset,
# batch_size=self.config.batch_size,
# shuffle=True
# )
# self.train_iter = mx.contrib.io.DataLoaderIter(self.train_dataloader)
else:
self.train_dataset = ImageFolderDataset(
root=config.image_dir,
flag=1,
transform=lambda data, label:
(data.astype(np.float32) / 255, label))
self.train_dataloader = mx.gluon.data.DataLoader(
dataset=self.train_dataset,
batch_size=self.config.batch_size,
shuffle=True)
self.eval_dataset = ImageFolderDataset(
root=config.image_dir,
flag=1,
transform=lambda data, label:
(data.astype(np.float32) / 255, label))
self.eval_dataloader = mx.gluon.data.DataLoader(
dataset=self.eval_dataset,
batch_size=self.config.batch_size,
shuffle=True)
def init_fit_data(self):
self.train_iter = ImageFolderDataIter(
root=self.config.image_dir,
data_shape=self.config.data_shape,
label_shape=(2, ),
data_names=['data'],
label_names=['softmax_label'],
flag=1,
transform=lambda data, label:
(data.astype(np.float32) / 255, label),
batch_size=self.config.batch_size)
def __init__(self, config):
self.config = config
self.ctxs = [mx.gpu(int(i)) for i in config.gpus.split(',')]
self.net = self.load_model_zoo(pretrained=self.config.pretrained_name,
epoch=self.config.pretrained_epoch)
self.net.collect_params().reset_ctx(self.ctxs)
self.trainer = mx.gluon.Trainer(self.net.collect_params(), 'sgd', {
'learning_rate': config.learning_rate,
'wd': config.weight_decay
})
if self.config.with_fit:
self.train_iter = ImageFolderDataIter(
root=self.config.image_dir,
data_shape=(3, 128, 128),
label_shape=(2, ),
data_names=['data'],
label_names=['softmax_label'],
flag=1,
transform=lambda data, label:
(data.astype(np.float32) / 255, label),
batch_size=self.config.batch_size)
# data_iter = iter(self.train_iter)
# batch = next(data_iter)
# logging.info(f'[] label: {batch.label}')
# # sys.exit(0)
# self.train_dataset = ImageFolderDataset(
# root=config.image_dir,
# flag=1,
# transform=lambda data, label: (data.astype(np.float32)/255, label)
# )
# self.train_dataloader = mx.gluon.data.DataLoader(
# dataset=self.train_dataset,
# batch_size=self.config.batch_size,
# shuffle=True
# )
# self.eval_dataset = ImageFolderDataset(
# root=config.image_dir,
# flag=1,
# transform=lambda data, label: (data.astype(np.float32)/255, label)
# )
# self.eval_dataloader = mx.gluon.data.DataLoader(
# dataset=self.eval_dataset,
# batch_size=self.config.batch_size,
# shuffle=True
# )
# self.train_iter = mx.contrib.io.DataLoaderIter(self.train_dataloader)
else:
self.train_dataset = ImageFolderDataset(
root=config.image_dir,
flag=1,
transform=lambda data, label:
(data.astype(np.float32) / 255, label))
self.train_dataloader = mx.gluon.data.DataLoader(
dataset=self.train_dataset,
batch_size=self.config.batch_size,
shuffle=True)
self.eval_dataset = ImageFolderDataset(
root=config.image_dir,
flag=1,
transform=lambda data, label:
(data.astype(np.float32) / 255, label))
self.eval_dataloader = mx.gluon.data.DataLoader(
dataset=self.eval_dataset,
batch_size=self.config.batch_size,
shuffle=True)
self.sec_loss = mx.gluon.loss.SoftmaxCrossEntropyLoss()
self.train_acc = mx.metric.Accuracy()
self.eval_acc = mx.metric.Accuracy()