796 lines
34 KiB
Python
796 lines
34 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import division
|
|
|
|
import copy
|
|
import os
|
|
from collections import defaultdict
|
|
from typing import Any, Callable, Dict, List, Tuple, Union, Optional
|
|
|
|
import numpy as np
|
|
import nvidia.dali.fn as fn
|
|
import nvidia.dali.ops as ops
|
|
import nvidia.dali.pipeline as pipeline
|
|
import nvidia.dali.types as types
|
|
import paddle
|
|
from nvidia.dali.plugin.paddle import DALIGenericIterator
|
|
from nvidia.dali.plugin.base_iterator import LastBatchPolicy
|
|
from ppcls.data.preprocess.ops.dali_operators import ColorJitter
|
|
from ppcls.data.preprocess.ops.dali_operators import CropImage
|
|
from ppcls.data.preprocess.ops.dali_operators import CropMirrorNormalize
|
|
from ppcls.data.preprocess.ops.dali_operators import DecodeImage
|
|
from ppcls.data.preprocess.ops.dali_operators import DecodeRandomResizedCrop
|
|
from ppcls.data.preprocess.ops.dali_operators import NormalizeImage
|
|
from ppcls.data.preprocess.ops.dali_operators import Pad
|
|
from ppcls.data.preprocess.ops.dali_operators import RandCropImage
|
|
from ppcls.data.preprocess.ops.dali_operators import RandCropImageV2
|
|
from ppcls.data.preprocess.ops.dali_operators import RandFlipImage
|
|
from ppcls.data.preprocess.ops.dali_operators import RandomCropImage
|
|
from ppcls.data.preprocess.ops.dali_operators import RandomRot90
|
|
from ppcls.data.preprocess.ops.dali_operators import RandomRotation
|
|
from ppcls.data.preprocess.ops.dali_operators import ResizeImage
|
|
from ppcls.data.preprocess.ops.dali_operators import ToCHWImage
|
|
from ppcls.engine.train.utils import type_name
|
|
from ppcls.utils import logger
|
|
|
|
INTERP_MAP = {
|
|
"nearest": types.DALIInterpType.INTERP_NN, # cv2.INTER_NEAREST
|
|
"bilinear": types.DALIInterpType.INTERP_LINEAR, # cv2.INTER_LINEAR
|
|
"bicubic": types.DALIInterpType.INTERP_CUBIC, # cv2.INTER_CUBIC
|
|
"lanczos": types.DALIInterpType.INTERP_LANCZOS3, # cv2.INTER_LANCZOS4
|
|
}
|
|
|
|
|
|
def make_pair(x: Union[Any, Tuple[Any], List[Any]]) -> Tuple[Any]:
|
|
"""repeat input x to be an tuple if x is an single element, else return x directly
|
|
|
|
Args:
|
|
x (Union[Any, Tuple[Any], List[Any]]): input x
|
|
|
|
Returns:
|
|
Tuple[Any]: tupled input
|
|
"""
|
|
return x if isinstance(x, (tuple, list)) else (x, x)
|
|
|
|
|
|
def parse_value_with_key(content: Union[Dict, List[Dict]],
|
|
key: str) -> Union[None, Any]:
|
|
"""parse value according to given key recursively, return None if not found
|
|
|
|
Args:
|
|
content (Union[Dict, List[Dict]]): content to be parsed
|
|
key (str): given key
|
|
|
|
Returns:
|
|
Union[None, Any]: result
|
|
"""
|
|
if isinstance(content, dict):
|
|
if key in content:
|
|
return content[key]
|
|
for content_ in content.values():
|
|
value = parse_value_with_key(content_, key)
|
|
if value is not None:
|
|
return value
|
|
elif isinstance(content, (tuple, list)):
|
|
for content_ in content:
|
|
value = parse_value_with_key(content_, key)
|
|
if value is not None:
|
|
return value
|
|
return None
|
|
|
|
|
|
def convert_cfg_to_dali(op_name: str, device: str, **op_cfg) -> Dict[str, Any]:
|
|
"""convert original preprocess op params into DALI-based op params
|
|
|
|
Args:
|
|
op_name (str): name of operator
|
|
device (str): device which operator applied on
|
|
|
|
Returns:
|
|
Dict[str, Any]: converted arguments for DALI initialization
|
|
"""
|
|
assert device in ["cpu", "gpu"
|
|
], f"device({device}) must in [\"cpu\", \"gpu\"]"
|
|
dali_op_cfg = {}
|
|
if op_name == "DecodeImage":
|
|
device = "cpu" if device == "cpu" else "mixed"
|
|
to_rgb = op_cfg.get("to_rgb", True)
|
|
channel_first = op_cfg.get("channel_first", False)
|
|
assert channel_first is False, \
|
|
f"`channel_first` must set to False when using DALI, but got {channel_first}"
|
|
dali_op_cfg.update({"device": device})
|
|
dali_op_cfg.update({
|
|
"output_type": types.DALIImageType.RGB
|
|
if to_rgb else types.DALIImageType.BGR
|
|
})
|
|
dali_op_cfg.update({
|
|
"device_memory_padding":
|
|
op_cfg.get("device_memory_padding", 211025920)
|
|
})
|
|
dali_op_cfg.update({
|
|
"host_memory_padding": op_cfg.get("host_memory_padding", 140544512)
|
|
})
|
|
elif op_name == "ResizeImage":
|
|
size = op_cfg.get("size", None)
|
|
resize_short = op_cfg.get("resize_short", None)
|
|
interpolation = op_cfg.get("interpolation", None)
|
|
if size is not None:
|
|
size = make_pair(size)
|
|
dali_op_cfg.update({"resize_y": size[0], "resize_x": size[1]})
|
|
if resize_short is not None:
|
|
dali_op_cfg.update({"resize_shorter": resize_short})
|
|
if interpolation is not None:
|
|
dali_op_cfg.update({"interp_type": INTERP_MAP[interpolation]})
|
|
elif op_name == "CropImage":
|
|
size = op_cfg.get("size", 224)
|
|
size = make_pair(size)
|
|
dali_op_cfg.update({"crop_h": size[1], "crop_w": size[0]})
|
|
dali_op_cfg.update({"crop_pos_x": 0.5, "crop_pos_y": 0.5})
|
|
elif op_name == "RandomCropImage":
|
|
size = op_cfg.get("size", 224)
|
|
if size is not None:
|
|
size = make_pair(size)
|
|
dali_op_cfg.update({"crop_h": size[1], "crop_w": size[0]})
|
|
elif op_name == "RandCropImage":
|
|
size = op_cfg.get("size", 224)
|
|
size = make_pair(size)
|
|
scale = op_cfg.get("scale", [0.08, 1.0])
|
|
ratio = op_cfg.get("ratio", [3.0 / 4, 4.0 / 3])
|
|
interpolation = op_cfg.get("interpolation", "bilinear")
|
|
dali_op_cfg.update({"size": size})
|
|
if scale is not None:
|
|
dali_op_cfg.update({"random_area": scale})
|
|
if ratio is not None:
|
|
dali_op_cfg.update({"random_aspect_ratio": ratio})
|
|
if interpolation is not None:
|
|
dali_op_cfg.update({"interp_type": INTERP_MAP[interpolation]})
|
|
elif op_name == "RandCropImageV2":
|
|
size = op_cfg.get("size", 224)
|
|
size = make_pair(size)
|
|
dali_op_cfg.update({"crop_h": size[1], "crop_w": size[0]})
|
|
elif op_name == "RandFlipImage":
|
|
prob = op_cfg.get("prob", 0.5)
|
|
flip_code = op_cfg.get("flip_code", 1)
|
|
dali_op_cfg.update({"prob": prob})
|
|
dali_op_cfg.update({"flip_code": flip_code})
|
|
elif op_name == "NormalizeImage":
|
|
# scale * (in - mean) / stddev + shift
|
|
scale = op_cfg.get("scale", 1.0 / 255.0)
|
|
if isinstance(scale, str):
|
|
scale = eval(scale)
|
|
mean = op_cfg.get("mean", [0.485, 0.456, 0.406])
|
|
std = op_cfg.get("std", [0.229, 0.224, 0.225])
|
|
mean = [v / scale for v in mean]
|
|
std = [v / scale for v in std]
|
|
order = op_cfg.get("order", "chw")
|
|
channel_num = op_cfg.get("channel_num", 3)
|
|
output_fp16 = op_cfg.get("output_fp16", False)
|
|
dali_op_cfg.update({
|
|
"mean": np.reshape(
|
|
np.array(
|
|
mean, dtype="float32"), [channel_num, 1, 1]
|
|
if order == "chw" else [1, 1, channel_num])
|
|
})
|
|
dali_op_cfg.update({
|
|
"stddev": np.reshape(
|
|
np.array(
|
|
std, dtype="float32"), [channel_num, 1, 1]
|
|
if order == "chw" else [1, 1, channel_num])
|
|
})
|
|
if output_fp16:
|
|
dali_op_cfg.update({"dtype": types.FLOAT16})
|
|
elif op_name == "ToCHWImage":
|
|
dali_op_cfg.update({"perm": [2, 0, 1]})
|
|
elif op_name == "ColorJitter":
|
|
prob = op_cfg.get("prob", 1.0)
|
|
brightness = op_cfg.get("brightness", 0.0)
|
|
contrast = op_cfg.get("contrast", 0.0)
|
|
saturation = op_cfg.get("saturation", 0.0)
|
|
hue = op_cfg.get("hue", 0.0)
|
|
dali_op_cfg.update({"prob": prob})
|
|
dali_op_cfg.update({"brightness_factor": brightness})
|
|
dali_op_cfg.update({"contrast_factor": contrast})
|
|
dali_op_cfg.update({"saturation_factor": saturation})
|
|
dali_op_cfg.update({"hue_factor": hue})
|
|
elif op_name == "RandomRotation":
|
|
prob = op_cfg.get("prob", 0.5)
|
|
degrees = op_cfg.get("degrees", 90)
|
|
interpolation = op_cfg.get("interpolation", "bilinear")
|
|
dali_op_cfg.update({"prob": prob})
|
|
dali_op_cfg.update({"angle": degrees})
|
|
dali_op_cfg.update({"interp_type": INTERP_MAP[interpolation]})
|
|
elif op_name == "Pad":
|
|
size = op_cfg.get("size", 224)
|
|
size = make_pair(size)
|
|
padding = op_cfg.get("padding", 0)
|
|
fill = op_cfg.get("fill", 0)
|
|
dali_op_cfg.update({
|
|
"crop_h": padding + size[1] + padding,
|
|
"crop_w": padding + size[0] + padding
|
|
})
|
|
dali_op_cfg.update({"fill_values": fill})
|
|
dali_op_cfg.update({"out_of_bounds_policy": "pad"})
|
|
elif op_name == "RandomRot90":
|
|
interpolation = op_cfg.get("interpolation", "nearest")
|
|
elif op_name == "DecodeRandomResizedCrop":
|
|
device = "cpu" if device == "cpu" else "mixed"
|
|
output_type = op_cfg.get("output_type", types.DALIImageType.RGB)
|
|
device_memory_padding = op_cfg.get("device_memory_padding", 211025920)
|
|
host_memory_padding = op_cfg.get("host_memory_padding", 140544512)
|
|
scale = op_cfg.get("scale", [0.08, 1.0])
|
|
ratio = op_cfg.get("ratio", [3.0 / 4, 4.0 / 3])
|
|
num_attempts = op_cfg.get("num_attempts", 100)
|
|
size = op_cfg.get("size", 224)
|
|
dali_op_cfg.update({"device": device})
|
|
if output_type is not None:
|
|
dali_op_cfg.update({"output_type": output_type})
|
|
if device_memory_padding is not None:
|
|
dali_op_cfg.update({
|
|
"device_memory_padding": device_memory_padding
|
|
})
|
|
if host_memory_padding is not None:
|
|
dali_op_cfg.update({"host_memory_padding": host_memory_padding})
|
|
if scale is not None:
|
|
dali_op_cfg.update({"random_area": scale})
|
|
if ratio is not None:
|
|
dali_op_cfg.update({"random_aspect_ratio": ratio})
|
|
if num_attempts is not None:
|
|
dali_op_cfg.update({"num_attempts": num_attempts})
|
|
if size is not None:
|
|
dali_op_cfg.update({"resize_x": size, "resize_y": size})
|
|
elif op_name == "CropMirrorNormalize":
|
|
dtype = types.FLOAT16 if op_cfg.get("output_fp16",
|
|
False) else types.FLOAT
|
|
output_layout = op_cfg.get("output_layout", "CHW")
|
|
size = op_cfg.get("size", None)
|
|
scale = op_cfg.get("scale", 1 / 255.0)
|
|
if isinstance(scale, str):
|
|
scale = eval(scale)
|
|
mean = op_cfg.get("mean", [0.485, 0.456, 0.406])
|
|
mean = [v / scale for v in mean]
|
|
std = op_cfg.get("std", [0.229, 0.224, 0.225])
|
|
std = [v / scale for v in std]
|
|
pad_output = op_cfg.get("channel_num", 3) == 4
|
|
prob = op_cfg.get("prob", 0.5)
|
|
dali_op_cfg.update({"dtype": dtype})
|
|
if output_layout is not None:
|
|
dali_op_cfg.update({"output_layout": output_layout})
|
|
if size is not None:
|
|
dali_op_cfg.update({"crop": (size, size)})
|
|
if mean is not None:
|
|
dali_op_cfg.update({"mean": mean})
|
|
if std is not None:
|
|
dali_op_cfg.update({"std": std})
|
|
if pad_output is not None:
|
|
dali_op_cfg.update({"pad_output": pad_output})
|
|
if prob is not None:
|
|
dali_op_cfg.update({"prob": prob})
|
|
else:
|
|
raise ValueError(
|
|
f"DALI operator \"{op_name}\" in PaddleClas is not implemented now. please refer to docs/zh_CN/training/config_description/develop_with_DALI.md"
|
|
)
|
|
if "device" not in dali_op_cfg:
|
|
dali_op_cfg.update({"device": device})
|
|
return dali_op_cfg
|
|
|
|
|
|
def build_dali_transforms(op_cfg_list: List[Dict[str, Any]],
|
|
mode: str,
|
|
device: str="gpu",
|
|
enable_fuse: bool=True) -> List[Callable]:
|
|
"""create dali operators based on the config
|
|
Args:
|
|
op_cfg_list (List[Dict[str, Any]]): a dict list, used to create some operators, such as config below
|
|
--------------------------------
|
|
- DecodeImage:
|
|
to_rgb: True
|
|
channel_first: False
|
|
- ResizeImage:
|
|
size: 224
|
|
- NormalizeImage:
|
|
scale: 0.00392157
|
|
mean: [0.485, 0.456, 0.406]
|
|
std: [0.229, 0.224, 0.225]
|
|
order: ""
|
|
--------------------------------
|
|
mode (str): mode.
|
|
device (str): device which dali operator(s) applied in. Defaults to "gpu".
|
|
enable_fuse (bool): whether to use fused dali operators instead of single operators, such as DecodeRandomResizedCrop. Defaults to True.
|
|
Returns:
|
|
List[Callable]: Callable DALI operators in list.
|
|
"""
|
|
assert isinstance(op_cfg_list, list), "operator config should be a list"
|
|
# build dali transforms list
|
|
|
|
dali_op_list = []
|
|
idx = 0
|
|
num_cfg_node = len(op_cfg_list)
|
|
while idx < num_cfg_node:
|
|
op_cfg = op_cfg_list[idx]
|
|
op_name = list(op_cfg)[0]
|
|
op_param = {} if op_cfg[op_name] is None else copy.deepcopy(op_cfg[
|
|
op_name])
|
|
fused_success = False
|
|
if enable_fuse:
|
|
# fuse operators if enabled
|
|
if idx + 1 < num_cfg_node:
|
|
op_name_nxt = list(op_cfg_list[idx + 1])[0]
|
|
if (op_name == "DecodeImage" and
|
|
op_name_nxt == "RandCropImage"):
|
|
fused_op_name = "DecodeRandomResizedCrop"
|
|
fused_op_param = convert_cfg_to_dali(
|
|
fused_op_name, device, **{
|
|
** op_param, ** (op_cfg_list[idx + 1][op_name_nxt])
|
|
})
|
|
fused_dali_op = eval(fused_op_name)(**fused_op_param)
|
|
idx += 2
|
|
dali_op_list.append(fused_dali_op)
|
|
fused_success = True
|
|
logger.info(
|
|
f"DALI fused Operator conversion({mode}): [DecodeImage, RandCropImage] -> {type_name(dali_op_list[-1])}: {fused_op_param}"
|
|
)
|
|
if not fused_success and 0 < idx and idx + 1 < num_cfg_node:
|
|
op_name_pre = list(op_cfg_list[idx - 1])[0]
|
|
op_name_nxt = list(op_cfg_list[idx + 1])[0]
|
|
if (op_name_pre == "RandCropImage" and
|
|
op_name == "RandFlipImage" and
|
|
op_name_nxt == "NormalizeImage"):
|
|
fused_op_name = "CropMirrorNormalize"
|
|
fused_op_param = convert_cfg_to_dali(
|
|
fused_op_name, device, **{
|
|
** op_param, **
|
|
(op_cfg_list[idx - 1][op_name_pre]), **
|
|
(op_cfg_list[idx + 1][op_name_nxt])
|
|
})
|
|
fused_dali_op = eval(fused_op_name)(**fused_op_param)
|
|
idx += 2
|
|
dali_op_list.append(fused_dali_op)
|
|
fused_success = True
|
|
logger.info(
|
|
f"DALI fused Operator conversion({mode}): [RandCropImage, RandFlipImage, NormalizeImage] -> {type_name(dali_op_list[-1])}: {fused_op_param}"
|
|
)
|
|
if not fused_success and idx + 1 < num_cfg_node:
|
|
op_name_nxt = list(op_cfg_list[idx + 1])[0]
|
|
if (op_name == "CropImage" and
|
|
op_name_nxt == "NormalizeImage"):
|
|
fused_op_name = "CropMirrorNormalize"
|
|
fused_op_param = convert_cfg_to_dali(
|
|
fused_op_name, device, **{
|
|
**
|
|
op_param,
|
|
**
|
|
(op_cfg_list[idx + 1][op_name_nxt]),
|
|
"prob": 0.0
|
|
})
|
|
fused_dali_op = eval(fused_op_name)(**fused_op_param)
|
|
idx += 2
|
|
dali_op_list.append(fused_dali_op)
|
|
fused_success = True
|
|
logger.info(
|
|
f"DALI fused Operator conversion({mode}): [CropImage, NormalizeImage] -> {type_name(dali_op_list[-1])}: {fused_op_param}"
|
|
)
|
|
if not enable_fuse or not fused_success:
|
|
assert isinstance(op_cfg,
|
|
dict) and len(op_cfg) == 1, "yaml format error"
|
|
if op_name == "Pad":
|
|
# NOTE: Argument `size` must be provided for DALI operator
|
|
op_param.update({
|
|
"size": parse_value_with_key(op_cfg_list[:idx], "size")
|
|
})
|
|
dali_param = convert_cfg_to_dali(op_name, device, **op_param)
|
|
dali_op = eval(op_name)(**dali_param)
|
|
dali_op_list.append(dali_op)
|
|
idx += 1
|
|
logger.info(
|
|
f"DALI Operator conversion({mode}): {op_name} -> {type_name(dali_op_list[-1])}: {dali_param}"
|
|
)
|
|
return dali_op_list
|
|
|
|
|
|
class ExternalSource_RandomIdentity(object):
|
|
"""PKsampler implemented with ExternalSource
|
|
|
|
Args:
|
|
batch_size (int): batch size
|
|
sample_per_id (int): number of instance(s) within an class
|
|
device_id (int): device id
|
|
shard_id (int): shard id
|
|
num_gpus (int): number of gpus
|
|
image_root (str): image root directory
|
|
cls_label_path (str): path to annotation file, such as `train_list.txt` or `val_list.txt`
|
|
delimiter (Optional[str], optional): delimiter. Defaults to None.
|
|
relabel (bool, optional): whether do relabel when original label do not starts from 0 or are discontinuous. Defaults to False.
|
|
sample_method (str, optional): sample method when generating prob_list. Defaults to "sample_avg_prob".
|
|
id_list (List[int], optional): list of (start_id, end_id, start_id, end_id) for set of ids to duplicated. Defaults to None.
|
|
ratio (List[Union[int, float]], optional): list of (ratio1, ratio2..) the duplication number for ids in id_list. Defaults to None.
|
|
shuffle (bool): whether to shuffle label list. Defaults to True.
|
|
"""
|
|
|
|
def __init__(self,
|
|
batch_size: int,
|
|
sample_per_id: int,
|
|
device_id: int,
|
|
shard_id: int,
|
|
num_gpus: int,
|
|
image_root: str,
|
|
cls_label_path: str,
|
|
delimiter: Optional[str]=None,
|
|
relabel: bool=False,
|
|
sample_method: str="sample_avg_prob",
|
|
id_list: List[int]=None,
|
|
ratio: List[Union[int, float]]=None,
|
|
shuffle: bool=True):
|
|
self.batch_size = batch_size
|
|
self.sample_per_id = sample_per_id
|
|
self.label_per_batch = self.batch_size // self.sample_per_id
|
|
self.device_id = device_id
|
|
self.shard_id = shard_id
|
|
self.num_gpus = num_gpus
|
|
self._img_root = image_root
|
|
self._cls_path = cls_label_path
|
|
self.delimiter = delimiter if delimiter is not None else " "
|
|
self.relabel = relabel
|
|
self.sample_method = sample_method
|
|
self.image_paths = []
|
|
self.labels = []
|
|
self.epoch = 0
|
|
|
|
# NOTE: code from ImageNetDataset below
|
|
with open(self._cls_path, "r") as fd:
|
|
lines = fd.readlines()
|
|
if self.relabel:
|
|
label_set = set()
|
|
for line in lines:
|
|
line = line.strip().split(self.delimiter)
|
|
label_set.add(np.int64(line[1]))
|
|
label_map = {
|
|
oldlabel: newlabel
|
|
for newlabel, oldlabel in enumerate(label_set)
|
|
}
|
|
|
|
for line in lines:
|
|
line = line.strip().split(self.delimiter)
|
|
self.image_paths.append(os.path.join(self._img_root, line[0]))
|
|
if self.relabel:
|
|
self.labels.append(label_map[np.int64(line[1])])
|
|
else:
|
|
self.labels.append(np.int64(line[1]))
|
|
assert os.path.exists(self.image_paths[
|
|
-1]), f"path {self.image_paths[-1]} does not exist."
|
|
|
|
# NOTE: code from PKSampler below
|
|
# group sample indexes into their label bucket
|
|
self.label_dict = defaultdict(list)
|
|
for idx, label in enumerate(self.labels):
|
|
self.label_dict[label].append(idx)
|
|
# get all label
|
|
self.label_list = list(self.label_dict)
|
|
assert len(self.label_list) * self.sample_per_id >= self.batch_size, \
|
|
f"batch size({self.batch_size}) should not be bigger than than #classes({len(self.label_list)})*sample_per_id({self.sample_per_id})"
|
|
|
|
if self.sample_method == "id_avg_prob":
|
|
self.prob_list = np.array([1 / len(self.label_list)] *
|
|
len(self.label_list))
|
|
elif self.sample_method == "sample_avg_prob":
|
|
counter = []
|
|
for label_i in self.label_list:
|
|
counter.append(len(self.label_dict[label_i]))
|
|
self.prob_list = np.array(counter) / sum(counter)
|
|
|
|
# reweight prob_list according to id_list and ratio if provided
|
|
if id_list and ratio:
|
|
assert len(id_list) % 2 == 0 and len(id_list) == len(ratio) * 2
|
|
for i in range(len(self.prob_list)):
|
|
for j in range(len(ratio)):
|
|
if i >= id_list[j * 2] and i <= id_list[j * 2 + 1]:
|
|
self.prob_list[i] = self.prob_list[i] * ratio[j]
|
|
break
|
|
self.prob_list = self.prob_list / sum(self.prob_list)
|
|
|
|
assert os.path.exists(
|
|
self._cls_path), f"path {self._cls_path} does not exist."
|
|
assert os.path.exists(
|
|
self._img_root), f"path {self._img_root} does not exist."
|
|
|
|
diff = np.abs(sum(self.prob_list) - 1)
|
|
if diff > 0.00000001:
|
|
self.prob_list[-1] = 1 - sum(self.prob_list[:-1])
|
|
if self.prob_list[-1] > 1 or self.prob_list[-1] < 0:
|
|
logger.error("PKSampler prob list error")
|
|
else:
|
|
logger.info(
|
|
"sum of prob list not equal to 1, diff is {}, change the last prob".
|
|
format(diff))
|
|
|
|
# whole dataset size
|
|
self.data_set_len = len(self.image_paths)
|
|
|
|
# get sharded size
|
|
self.sharded_data_set_len = self.data_set_len // self.num_gpus
|
|
|
|
# iteration log
|
|
self.shuffle = shuffle
|
|
self.total_iter = self.sharded_data_set_len // batch_size
|
|
self.iter_count = 0
|
|
|
|
def __iter__(self):
|
|
if self.shuffle:
|
|
seed = self.shard_id * 12345 + self.epoch
|
|
np.random.RandomState(seed).shuffle(self.label_list)
|
|
np.random.RandomState(seed).shuffle(self.prob_list)
|
|
self.epoch += 1
|
|
return self
|
|
|
|
def __next__(self):
|
|
if self.iter_count >= self.total_iter:
|
|
self.__iter__()
|
|
self.iter_count = 0
|
|
|
|
batch_indexes = []
|
|
for _ in range(self.sharded_data_set_len):
|
|
batch_label_list = np.random.choice(
|
|
self.label_list,
|
|
size=self.label_per_batch,
|
|
replace=False,
|
|
p=self.prob_list)
|
|
for label_i in batch_label_list:
|
|
label_i_indexes = self.label_dict[label_i]
|
|
if self.sample_per_id <= len(label_i_indexes):
|
|
batch_indexes.extend(
|
|
np.random.choice(
|
|
label_i_indexes,
|
|
size=self.sample_per_id,
|
|
replace=False))
|
|
else:
|
|
batch_indexes.extend(
|
|
np.random.choice(
|
|
label_i_indexes,
|
|
size=self.sample_per_id,
|
|
replace=True))
|
|
if len(batch_indexes) == self.batch_size:
|
|
break
|
|
batch_indexes = []
|
|
|
|
batch_raw_images = []
|
|
batch_labels = []
|
|
for index in batch_indexes:
|
|
batch_raw_images.append(
|
|
np.fromfile(
|
|
self.image_paths[index], dtype="uint8"))
|
|
batch_labels.append(self.labels[index])
|
|
|
|
self.iter_count += 1
|
|
return (batch_raw_images, np.array(batch_labels, dtype="int64"))
|
|
|
|
def __len__(self):
|
|
return self.sharded_data_set_len
|
|
|
|
|
|
class HybridPipeline(pipeline.Pipeline):
|
|
"""Hybrid Pipeline
|
|
|
|
Args:
|
|
device (str): device
|
|
batch_size (int): batch size
|
|
py_num_workers (int): number of python worker(s)
|
|
num_threads (int): number of thread(s)
|
|
device_id (int): device id
|
|
seed (int): random seed
|
|
file_root (str): file root path
|
|
file_list (str): path to annotation file, such as `train_list.txt` or `val_list.txt`
|
|
transform_list (List[Callable]): List of DALI transform operator(s)
|
|
shard_id (int, optional): shard id. Defaults to 0.
|
|
num_shards (int, optional): number of shard(s). Defaults to 1.
|
|
random_shuffle (bool, optional): whether shuffle data during training. Defaults to True.
|
|
ext_src (optional): custom external source. Defaults to None.
|
|
"""
|
|
|
|
def __init__(self,
|
|
device: str,
|
|
batch_size: int,
|
|
py_num_workers: int,
|
|
num_threads: int,
|
|
device_id: int,
|
|
seed: int,
|
|
file_root: str,
|
|
file_list: str,
|
|
transform_list: List[Callable],
|
|
shard_id: int=0,
|
|
num_shards: int=1,
|
|
random_shuffle: bool=True,
|
|
ext_src=None):
|
|
super(HybridPipeline, self).__init__(
|
|
batch_size=batch_size,
|
|
device_id=device_id,
|
|
seed=seed,
|
|
py_start_method="fork" if ext_src is None else "spawn",
|
|
py_num_workers=py_num_workers,
|
|
num_threads=num_threads)
|
|
self.device = device
|
|
self.ext_src = ext_src
|
|
if ext_src is None:
|
|
self.reader = ops.readers.File(
|
|
file_root=file_root,
|
|
file_list=file_list,
|
|
shard_id=shard_id,
|
|
num_shards=num_shards,
|
|
random_shuffle=random_shuffle)
|
|
self.transforms = ops.Compose(transform_list)
|
|
self.cast = ops.Cast(dtype=types.DALIDataType.INT64, device=device)
|
|
|
|
def define_graph(self):
|
|
if self.ext_src:
|
|
raw_images, labels = fn.external_source(
|
|
source=self.ext_src,
|
|
num_outputs=2,
|
|
dtype=[types.DALIDataType.UINT8, types.DALIDataType.INT64],
|
|
batch=True,
|
|
parallel=True)
|
|
else:
|
|
raw_images, labels = self.reader(name="Reader")
|
|
images = self.transforms(raw_images)
|
|
return [
|
|
images, self.cast(labels.gpu() if self.device == "gpu" else labels)
|
|
]
|
|
|
|
def __len__(self):
|
|
if self.ext_src is not None:
|
|
return len(self.ext_src)
|
|
return self.epoch_size(name="Reader")
|
|
|
|
|
|
class DALIImageNetIterator(DALIGenericIterator):
|
|
def __init__(self, *kargs, **kwargs):
|
|
super(DALIImageNetIterator, self).__init__(*kargs, **kwargs)
|
|
self.in_dynamic_mode = paddle.in_dynamic_mode()
|
|
|
|
def __next__(self) -> List[paddle.Tensor]:
|
|
data_batch = super(DALIImageNetIterator,
|
|
self).__next__() # List[Dict[str, Tensor], ...]
|
|
# reformat to List[Tensor1, Tensor2, ...]
|
|
data_batch = [
|
|
paddle.to_tensor(data_batch[0][key])
|
|
if self.in_dynamic_mode else data_batch[0][key]
|
|
for key in self.output_map
|
|
]
|
|
return data_batch
|
|
|
|
|
|
def dali_dataloader(config: Dict[str, Any],
|
|
mode: str,
|
|
device: str,
|
|
py_num_workers: int=1,
|
|
num_threads: int=4,
|
|
seed: Optional[int]=None,
|
|
enable_fuse: bool=True) -> DALIImageNetIterator:
|
|
"""build and return HybridPipeline
|
|
|
|
Args:
|
|
config (Dict[str, Any]): train/eval dataloader configuration
|
|
mode (str): mode
|
|
device (str): device string
|
|
py_num_workers (int, optional): number of python worker(s). Defaults to 1.
|
|
num_threads (int, optional): number of thread(s). Defaults to 4.
|
|
seed (Optional[int], optional): random seed. Defaults to None.
|
|
enable_fuse (bool, optional): enable fused operator(s). Defaults to True.
|
|
|
|
Returns:
|
|
DALIImageNetIterator: Iterable DALI dataloader
|
|
"""
|
|
assert "gpu" in device, f"device must be \"gpu\" when running with DALI, but got {device}"
|
|
config_dataloader = config[mode]
|
|
device_id = int(device.split(":")[1])
|
|
device = "gpu"
|
|
seed = 42 if seed is None else seed
|
|
env = os.environ
|
|
num_gpus = paddle.distributed.get_world_size()
|
|
|
|
batch_size = config_dataloader["sampler"]["batch_size"]
|
|
file_root = config_dataloader["dataset"]["image_root"]
|
|
file_list = config_dataloader["dataset"]["cls_label_path"]
|
|
sampler_name = config_dataloader["sampler"].get("name",
|
|
"DistributedBatchSampler")
|
|
transform_ops_cfg = config_dataloader["dataset"]["transform_ops"]
|
|
random_shuffle = config_dataloader["sampler"].get("shuffle", None)
|
|
dali_transforms = build_dali_transforms(
|
|
transform_ops_cfg, mode, device, enable_fuse=enable_fuse)
|
|
if "ToCHWImage" not in [type_name(op) for op in dali_transforms] and (
|
|
"CropMirrorNormalize" not in
|
|
[type_name(op) for op in dali_transforms]):
|
|
dali_transforms.append(ToCHWImage(perm=[2, 0, 1], device=device))
|
|
logger.info(
|
|
"Append DALI operator \"ToCHWImage\" at the end of dali_transforms for getting output in \"CHW\" shape"
|
|
)
|
|
|
|
if mode.lower() in ["train"]:
|
|
if "PADDLE_TRAINER_ID" in env and "PADDLE_TRAINERS_NUM" in env and "FLAGS_selected_gpus" in env:
|
|
shard_id = int(env["PADDLE_TRAINER_ID"])
|
|
num_shards = int(env["PADDLE_TRAINERS_NUM"])
|
|
device_id = int(env["FLAGS_selected_gpus"])
|
|
else:
|
|
shard_id = 0
|
|
num_shards = 1
|
|
logger.info(
|
|
f"Building DALI {mode} pipeline with num_shards: {num_shards}, num_gpus: {num_gpus}"
|
|
)
|
|
|
|
random_shuffle = random_shuffle if random_shuffle is not None else True
|
|
if sampler_name in ["PKSampler", "DistributedRandomIdentitySampler"]:
|
|
ext_src = ExternalSource_RandomIdentity(
|
|
batch_size=batch_size,
|
|
sample_per_id=config_dataloader["sampler"][
|
|
"sample_per_id"
|
|
if sampler_name == "PKSampler" else "num_instances"],
|
|
device_id=device_id,
|
|
shard_id=shard_id,
|
|
num_gpus=num_gpus,
|
|
image_root=file_root,
|
|
cls_label_path=file_list,
|
|
delimiter=None,
|
|
relabel=config_dataloader["dataset"].get("relabel", False),
|
|
sample_method=config_dataloader["sampler"].get(
|
|
"sample_method", "sample_avg_prob"),
|
|
id_list=config_dataloader["sampler"].get("id_list", None),
|
|
ratio=config_dataloader["sampler"].get("ratio", None),
|
|
shuffle=random_shuffle)
|
|
logger.info(
|
|
f"Building DALI {mode} pipeline with ext_src({type_name(ext_src)})"
|
|
)
|
|
else:
|
|
ext_src = None
|
|
|
|
pipe = HybridPipeline(device, batch_size, py_num_workers, num_threads,
|
|
device_id, seed + shard_id, file_root, file_list,
|
|
dali_transforms, shard_id, num_shards,
|
|
random_shuffle, ext_src)
|
|
pipe.build()
|
|
pipelines = [pipe]
|
|
if ext_src is None:
|
|
return DALIImageNetIterator(
|
|
pipelines, ["data", "label"], reader_name="Reader")
|
|
else:
|
|
return DALIImageNetIterator(
|
|
pipelines,
|
|
["data", "label"],
|
|
size=len(ext_src),
|
|
last_batch_policy=LastBatchPolicy.
|
|
DROP # make reset() successfully
|
|
)
|
|
elif mode.lower() in ["eval", "gallery", "query"]:
|
|
assert sampler_name in ["DistributedBatchSampler"], \
|
|
f"sampler_name({sampler_name}) must in [\"DistributedBatchSampler\"]"
|
|
if "PADDLE_TRAINER_ID" in env and "PADDLE_TRAINERS_NUM" in env and "FLAGS_selected_gpus" in env:
|
|
shard_id = int(env["PADDLE_TRAINER_ID"])
|
|
num_shards = int(env["PADDLE_TRAINERS_NUM"])
|
|
device_id = int(env["FLAGS_selected_gpus"])
|
|
else:
|
|
shard_id = 0
|
|
num_shards = 1
|
|
logger.info(
|
|
f"Building DALI {mode} pipeline with num_shards: {num_shards}, num_gpus: {num_gpus}..."
|
|
)
|
|
|
|
random_shuffle = random_shuffle if random_shuffle is not None else False
|
|
pipe = HybridPipeline(device, batch_size, py_num_workers, num_threads,
|
|
device_id, seed + shard_id, file_root, file_list,
|
|
dali_transforms, shard_id, num_shards,
|
|
random_shuffle)
|
|
pipe.build()
|
|
pipelines = [pipe]
|
|
return DALIImageNetIterator(
|
|
pipelines, ["data", "label"], reader_name="Reader")
|
|
else:
|
|
raise ValueError(f"Invalid mode({mode}) when building DALI pipeline")
|