Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add PoolDataLoader for parrots #134

Merged
merged 7 commits into from
Sep 3, 2020
Merged

add PoolDataLoader for parrots #134

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
4a4e5c6
add PoolDataLoader for parrots
magicdream2222 Sep 1, 2020
b9f6916
fix lint
magicdream2222 Sep 1, 2020
348ef57
fix lint
magicdream2222 Sep 1, 2020
18c2c9a
pin_memory and prefetch_num can be got by config
magicdream2222 Sep 2, 2020
206b90a
fix lint
magicdream2222 Sep 2, 2020
8a84dd3
fix lint
magicdream2222 Sep 2, 2020
85ee4e5
update train.py
magicdream2222 Sep 3, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 79 additions & 31 deletions mmedit/apis/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,15 +99,29 @@ def _dist_train(model,
"""
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
dist=True,
drop_last=cfg.data.get('drop_last', False),
seed=cfg.seed) for ds in dataset
]
if torch.__version__ == 'parrots':
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
dist=True,
drop_last=cfg.data.get('drop_last', False),
seed=cfg.seed,
prefetch_num=cfg.data.get('prefetch_num', 2),
pin_memory=cfg.data.get('pin_memory', False)) for ds in dataset
]
else:
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
dist=True,
drop_last=cfg.data.get('drop_last', False),
seed=cfg.seed) for ds in dataset
]

# put model on gpus
find_unused_parameters = cfg.get('find_unused_parameters', False)
model = DistributedDataParallelWrapper(
Expand Down Expand Up @@ -146,12 +160,22 @@ def _dist_train(model,
cfg.data.samples_per_gpu)
workers_per_gpu = cfg.data.get('val_workers_per_gpu',
cfg.data.workers_per_gpu)
data_loader = build_dataloader(
dataset,
samples_per_gpu=samples_per_gpu,
workers_per_gpu=workers_per_gpu,
dist=True,
shuffle=False)
if torch.__version__ == 'parrots':
data_loader = build_dataloader(
dataset,
samples_per_gpu=samples_per_gpu,
workers_per_gpu=workers_per_gpu,
dist=True,
shuffle=False,
prefetch_num=cfg.data.get('prefetch_num', 2),
pin_memory=cfg.data.get('pin_memory', False))
else:
data_loader = build_dataloader(
dataset,
samples_per_gpu=samples_per_gpu,
workers_per_gpu=workers_per_gpu,
dist=True,
shuffle=False)
save_path = osp.join(cfg.work_dir, 'val_visuals')
runner.register_hook(
DistEvalIterHook(
Expand Down Expand Up @@ -185,16 +209,30 @@ def _non_dist_train(model,
"""
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
cfg.gpus,
dist=False,
drop_last=cfg.data.get('drop_last', False),
seed=cfg.seed) for ds in dataset
]
if torch.__version__ == 'parrots':
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
cfg.gpus,
dist=False,
drop_last=cfg.data.get('drop_last', False),
seed=cfg.seed,
prefetch_num=cfg.data.get('prefetch_num', 2),
pin_memory=cfg.data.get('pin_memory', False)) for ds in dataset
]
else:
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
cfg.gpus,
dist=False,
drop_last=cfg.data.get('drop_last', False),
seed=cfg.seed) for ds in dataset
]
# put model on gpus
model = MMDataParallel(model, device_ids=range(cfg.gpus)).cuda()

Expand Down Expand Up @@ -229,12 +267,22 @@ def _non_dist_train(model,
cfg.data.samples_per_gpu)
workers_per_gpu = cfg.data.get('val_workers_per_gpu',
cfg.data.workers_per_gpu)
data_loader = build_dataloader(
dataset,
samples_per_gpu=samples_per_gpu,
workers_per_gpu=workers_per_gpu,
dist=True,
shuffle=False)
if torch.__version__ == 'parrots':
data_loader = build_dataloader(
dataset,
samples_per_gpu=samples_per_gpu,
workers_per_gpu=workers_per_gpu,
dist=True,
shuffle=False,
prefetch_num=cfg.data.get('prefetch_num', 2),
pin_memory=cfg.data.get('pin_memory', False))
else:
data_loader = build_dataloader(
dataset,
samples_per_gpu=samples_per_gpu,
workers_per_gpu=workers_per_gpu,
dist=True,
shuffle=False)
save_path = osp.join(cfg.work_dir, 'val_visuals')
runner.register_hook(
EvalIterHook(data_loader, save_path=save_path, **cfg.evaluation))
Expand Down
8 changes: 7 additions & 1 deletion mmedit/datasets/builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,21 @@
from functools import partial

import numpy as np
import torch
from mmcv.parallel import collate
from mmcv.runner import get_dist_info
from mmcv.utils import build_from_cfg
from torch.utils.data import ConcatDataset, DataLoader
from torch.utils.data import ConcatDataset

from .dataset_wrappers import RepeatDataset
from .registry import DATASETS
from .samplers import DistributedSampler

if torch.__version__ == 'parrots':
from torch.utils.data import PoolDataLoader as DataLoader
else:
from torch.utils.data import DataLoader

if platform.system() != 'Windows':
# https://github.com/pytorch/pytorch/issues/973
import resource
Expand Down