def get_test_dataset(cfg: DictConfig):
"""
Get test dataset
Args:
cfg:
Returns:
"""
test_img_dir = f'{cfg.data.folder_path}/test'
valid_augs_list = [
load_obj(i['class_name'])(**i['params'])
for i in cfg['augmentation']['valid']['augs']
]
valid_bbox_params = OmegaConf.to_container(
(cfg['augmentation']['valid']['bbox_params']))
valid_augs = A.Compose(valid_augs_list, bbox_params=valid_bbox_params)
dataset_class = load_obj(cfg.dataset.class_name)
test_dataset = dataset_class(None, 'test', test_img_dir, cfg, valid_augs)
return test_dataset
def configure_optimizers(self):
if 'decoder_lr' in self.cfg.optimizer.params.keys():
params = [
{
'params': self.model.decoder.parameters(),
'lr': self.cfg.optimizer.params.lr
},
{
'params': self.model.encoder.parameters(),
'lr': self.cfg.optimizer.params.decoder_lr
},
]
optimizer = load_obj(self.cfg.optimizer.class_name)(params)
else:
optimizer = load_obj(self.cfg.optimizer.class_name)(
self.model.parameters(), **self.cfg.optimizer.params)
scheduler = load_obj(self.cfg.scheduler.class_name)(
optimizer, **self.cfg.scheduler.params)
return [optimizer], [{
'scheduler': scheduler,
'interval': self.cfg.scheduler.step,
'monitor': self.cfg.scheduler.monitor
}]
def get_m5model(cfg: DictConfig) -> Any:
"""
Get model
Args:
cfg: config
Returns:
initialized model
"""
backcast_length = cfg.dataset.backcast_length
forecast_length = cfg.dataset.forecast_length
f_b_dim = (forecast_length, backcast_length)
# collect stack parameters.
model_dict: Dict[str, list] = {k: [] for k in cfg.model.blocks[0].keys()}
for block in cfg.model.blocks:
for k, v in block.items():
if type(v) == str:
v = load_obj(v)
model_dict[k].append(v)
criterion = load_obj(cfg.loss.class_name)(**cfg.loss.params)
net = load_obj(cfg.model.class_name)
net = net(f_b_dim=f_b_dim, criterion=criterion, **model_dict)
return net
def configure_optimizers(self):
optimizer = load_obj(self.cfg.optimizer.class_name)(self.net.parameters(), **self.cfg.optimizer.params)
scheduler = load_obj(self.cfg.scheduler.class_name)(optimizer, **self.cfg.scheduler.params)
return (
[optimizer],
[{'scheduler': scheduler, 'interval': self.cfg.scheduler.step, 'monitor': self.cfg.scheduler.monitor}],
)
def __init__(self, cfg: DictConfig) -> None:
"""
Model class.
Args:
cfg: main config
"""
super().__init__()
self.encoder = load_obj(
cfg.model.encoder.class_name)(**cfg.model.encoder.params)
self.decoder = load_obj(cfg.model.decoder.class_name)(
output_dimension=self.encoder.output_dimension,
**cfg.model.decoder.params)
self.loss = load_obj(cfg.loss.class_name)()
def get_datasets(cfg: DictConfig) -> dict:
"""
Get datases for modelling
Args:
cfg: config
Returns:
dict with datasets
"""
train_df = pd.read_csv(
f'{cfg.data.folder_path}/data/sales_train_evaluation.csv')
train_fold_df = train_df.iloc[:, :-28]
valid_fold_df = train_df.iloc[:, -28:].copy()
backcast_length = cfg.dataset.backcast_length
# forecast_length = cfg.dataset.forecast_length
# train dataset
dataset_class = load_obj(cfg.dataset.class_name)
# get only useful columns
train_data = pd.concat(
[train_fold_df.iloc[:, 0], train_fold_df.iloc[:, 6:-1]], axis=1)
train_dataset = dataset_class(df=train_data, mode='train', cfg=cfg)
valid_data = pd.concat([
train_fold_df.iloc[:, 0],
train_fold_df.iloc[:, -backcast_length - 1:-1], valid_fold_df
],
axis=1)
valid_dataset = dataset_class(df=valid_data, mode='valid', cfg=cfg)
return {'train': train_dataset, 'valid': valid_dataset}
def __init__(self, hparams: Dict[str, float], cfg: DictConfig):
super(LitM5NBeats, self).__init__()
self.cfg = cfg
self.hparams: Dict[str, float] = hparams
self.net = get_m5model(self.cfg)
self.hparams['n_params'] = sum(p.numel() for p in self.net.parameters())
self.criterion = load_obj(self.cfg.loss.class_name)(**self.cfg.loss.params)
self._prepare_evaluator()
def __init__(self, cfg: DictConfig):
super().__init__()
self.model = load_obj(cfg.model.backbone.class_name)
self.model = self.model(**cfg.model.backbone.params)
# get number of input features for the classifier
self.output_dimension = self.model.roi_heads.box_predictor.cls_score.in_features
def get_training_datasets(cfg: DictConfig) -> Dict:
"""
Get datases for modelling
Args:
cfg: config
Returns:
dict with datasets
"""
train = pd.read_csv(f'{cfg.data.folder_path}/train.csv')
train[['x', 'y', 'w', 'h']] = pd.DataFrame(
np.stack(train['bbox'].apply(lambda x: ast.literal_eval(x)))).astype(
np.float32)
# precalculate some values
train['x1'] = train['x'] + train['w']
train['y1'] = train['y'] + train['h']
train['area'] = train['w'] * train['h']
train_ids, valid_ids = train_test_split(train['image_id'].unique(),
test_size=0.1,
random_state=cfg.training.seed)
# for fast training
if cfg.training.debug:
train_ids = train_ids[:10]
valid_ids = valid_ids[:10]
train_df = train.loc[train['image_id'].isin(train_ids)]
valid_df = train.loc[train['image_id'].isin(valid_ids)]
train_img_dir = f'{cfg.data.folder_path}/train'
# train dataset
dataset_class = load_obj(cfg.dataset.class_name)
# initialize augmentations
train_augs = load_augs(cfg['augmentation']['train']['augs'])
valid_augs = load_augs(cfg['augmentation']['valid']['augs'])
train_dataset = dataset_class(dataframe=train_df,
mode='train',
image_dir=train_img_dir,
cfg=cfg,
transforms=train_augs)
valid_dataset = dataset_class(dataframe=valid_df,
mode='valid',
image_dir=train_img_dir,
cfg=cfg,
transforms=valid_augs)
return {'train': train_dataset, 'valid': valid_dataset}
def get_training_datasets(cfg: DictConfig) -> dict:
data_dir = f'{cfg.data.folder_path}'
data_ids = [
fname.split('.')[0] for fname in os.listdir(f'{data_dir}/img_train')
]
train_ids, valid_ids = train_test_split(data_ids,
test_size=0.1,
random_state=cfg.training.seed)
# for fast training
if cfg.training.debug:
train_ids = train_ids[:10]
valid_ids = valid_ids[:10]
# dataset
dataset_class = load_obj(cfg.dataset.class_name)
# initialize augmentations
train_augs_list = [
load_obj(i['class_name'])(**i['params'])
for i in cfg['augmentation']['train']['augs']
]
train_augs = A.Compose(train_augs_list)
valid_augs_list = [
load_obj(i['class_name'])(**i['params'])
for i in cfg['augmentation']['valid']['augs']
]
valid_augs = A.Compose(valid_augs_list)
train_dataset = dataset_class(train_ids, 'train', data_dir, cfg,
train_augs)
valid_dataset = dataset_class(valid_ids, 'valid', data_dir, cfg,
valid_augs)
return {'train': train_dataset, 'valid': valid_dataset}
def get_wheat_model(cfg: DictConfig = None):
"""
Get model
Args:
cfg: config
Returns:
"""
model = load_obj(cfg.model.backbone.class_name)
model = model(**cfg.model.backbone.params)
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
head = load_obj(cfg.model.head.class_name)
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = head(in_features, cfg.model.head.params.num_classes)
return model
def get_unet_model(cfg: DictConfig = None):
"""
Get model
Args:
cfg: config
Returns:
"""
model = load_obj(cfg.model.backbone.class_name)
model = model(**cfg.model.backbone.params)
return model
def get_test_dataset(cfg: DictConfig):
"""
Get test dataset
Args:
cfg:
Returns:
"""
test_img_dir = f'{cfg.data.folder_path}/data55401/img_testA'
valid_augs_list = [
load_obj(i['class_name'])(**i['params'])
for i in cfg['augmentation']['valid']['augs']
]
valid_augs = A.Compose(valid_augs_list)
dataset_class = load_obj(cfg.dataset.class_name)
test_dataset = dataset_class(None, 'test', test_img_dir, cfg, valid_augs)
return test_dataset
def load_augs(cfg: DictConfig) -> A.Compose:
"""
Load albumentations
Args:
cfg:
Returns:
compose object
"""
augs = []
for a in cfg:
if a['class_name'] == 'albumentations.OneOf':
small_augs = []
for small_aug in a['params']:
# yaml can't contain tuples, so we need to convert manually
params = {
k: (v if type(v) != omegaconf.listconfig.ListConfig else
tuple(v))
for k, v in small_aug['params'].items()
}
aug = load_obj(small_aug['class_name'])(**params)
small_augs.append(aug)
aug = load_obj(a['class_name'])(small_augs)
augs.append(aug)
else:
params = {
k:
(v if type(v) != omegaconf.listconfig.ListConfig else tuple(v))
for k, v in a['params'].items()
}
aug = load_obj(a['class_name'])(**params)
augs.append(aug)
return A.Compose(augs)
def get_test_dataset(cfg: DictConfig) -> object:
"""
Get test dataset
Args:
cfg:
Returns:
test dataset
"""
test_img_dir = f'{cfg.data.folder_path}/test'
valid_augs = load_augs(cfg['augmentation']['valid']['augs'])
dataset_class = load_obj(cfg.dataset.class_name)
test_dataset = dataset_class(dataframe=None,
mode='test',
image_dir=test_img_dir,
cfg=cfg,
transforms=valid_augs)
return test_dataset
def get_training_datasets(cfg: DictConfig) -> dict:
"""
Get datases for modelling
Args:
cfg: config
Returns:
"""
train = pd.read_csv(f'{cfg.data.folder_path}/train.csv')
train[['x', 'y', 'w', 'h']] = pd.DataFrame(
np.stack(train['bbox'].apply(lambda x: ast.literal_eval(x)))).astype(
np.float32)
# precalculate some values
train['x1'] = train['x'] + train['w']
train['y1'] = train['y'] + train['h']
train['area'] = train['w'] * train['h']
train_ids, valid_ids = train_test_split(train['image_id'].unique(),
test_size=0.1,
random_state=cfg.training.seed)
# for fast training
if cfg.training.debug:
train_ids = train_ids[:10]
valid_ids = valid_ids[:10]
train_df = train.loc[train['image_id'].isin(train_ids)]
valid_df = train.loc[train['image_id'].isin(valid_ids)]
train_img_dir = f'{cfg.data.folder_path}/train'
# train dataset
dataset_class = load_obj(cfg.dataset.class_name)
# initialize augmentations
train_augs_list = [
load_obj(i['class_name'])(**i['params'])
for i in cfg['augmentation']['train']['augs']
]
train_bbox_params = OmegaConf.to_container(
(cfg['augmentation']['train']['bbox_params']))
train_augs = A.Compose(train_augs_list, bbox_params=train_bbox_params)
valid_augs_list = [
load_obj(i['class_name'])(**i['params'])
for i in cfg['augmentation']['valid']['augs']
]
valid_bbox_params = OmegaConf.to_container(
(cfg['augmentation']['valid']['bbox_params']))
valid_augs = A.Compose(valid_augs_list, bbox_params=valid_bbox_params)
train_dataset = dataset_class(train_df, 'train', train_img_dir, cfg,
train_augs)
valid_dataset = dataset_class(valid_df, 'valid', train_img_dir, cfg,
valid_augs)
return {'train': train_dataset, 'valid': valid_dataset}
def configure_optimizers(self):
optimizer = load_obj(self.hparams.optimizer.class_name)(self.model.parameters(),
**self.hparams.optimizer.params)
scheduler = load_obj(self.hparams.scheduler.class_name)(optimizer, **self.hparams.scheduler.params)
return [optimizer], [scheduler]
