class EvaluationSettings(Configurable):
data_loaders = State()
visualize = State(default=True)
resume = State()
def __init__(self, **kwargs):
self.load_all(**kwargs)
class RandomSampleDataLoader(Configurable, torch.utils.data.DataLoader):
datasets = State()
weights = State()
batch_size = State(default=256)
num_workers = State(default=10)
size = State(default=2**31)
def __init__(self, **kwargs):
self.load_all(**kwargs)
cmd = kwargs['cmd']
if 'batch_size' in cmd:
self.batch_size = cmd['batch_size']
probs = []
for dataset, weight in zip(self.datasets, self.weights):
probs.append(np.full(len(dataset), weight / len(dataset)))
dataset = ConcatDataset(self.datasets)
probs = np.concatenate(probs)
assert (len(dataset) == len(probs))
sampler = RandomSampleSampler(dataset, probs, self.size)
torch.utils.data.DataLoader.__init__(
self,
dataset,
batch_size=self.batch_size,
num_workers=self.num_workers,
sampler=sampler,
worker_init_fn=default_worker_init_fn,
)
class ResizeData(_ResizeImage, DataProcess):
key = State(default='image')
box_key = State(default='polygons')
image_size = State(default=[64, 256]) # height, width
def __init__(self, cmd={}, mode=None, key=None, box_key=None, **kwargs):
self.load_all(**kwargs)
if mode is not None:
self.mode = mode
if key is not None:
self.key = key
if box_key is not None:
self.box_key = box_key
if 'resize_mode' in cmd:
self.mode = cmd['resize_mode']
assert self.mode in self.MODES
def process(self, data):
height, width = data['image'].shape[:2]
new_height, new_width = self.get_image_size(data['image'])
data[self.key] = self.resize_or_pad(data[self.key])
charboxes = data[self.box_key]
data[self.box_key] = charboxes.copy()
data[self.box_key][:, :, 0] = data[self.box_key][:, :, 0] * \
new_width / width
data[self.box_key][:, :, 1] = data[self.box_key][:, :, 1] * \
new_height / height
return data
class ShowSettings(Configurable):
data_loader = State()
representer = State()
visualizer = State()
def __init__(self, **kwargs):
self.load_all(**kwargs)
class Checkpoint(Configurable):
start_epoch = State(default=0)
start_iter = State(default=0)
resume = State()
def __init__(self, **kwargs):
self.load_all(**kwargs)
cmd = kwargs['cmd']
if 'start_epoch' in cmd:
self.start_epoch = cmd['start_epoch']
if 'start_iter' in cmd:
self.start_iter = cmd['start_iter']
if 'resume' in cmd:
self.resume = cmd['resume']
def restore_model(self, model, device, logger):
if self.resume is None:
return
if not os.path.exists(self.resume):
self.logger.warning("Checkpoint not found: " + self.resume)
return
logger.info("Resuming from " + self.resume)
state_dict = torch.load(self.resume, map_location=device)
model.load_state_dict(state_dict, strict=False)
logger.info("Resumed from " + self.resume)
def restore_counter(self):
return self.start_epoch, self.start_iter
class FilterKeys(DataProcess):
required = State(default=[])
superfluous = State(default=[])
def __init__(self, **kwargs):
super().__init__(self, **kwargs)
self.required_keys = set(self.required)
self.superfluous_keys = set(self.superfluous)
if len(self.required_keys) > 0 and len(self.superfluous_keys) > 0:
raise ValueError(
'required_keys and superfluous_keys can not be specified at the same time.')
def process(self, data):
for key in self.required:
assert key in data, '%s is required in data' % key
superfluous = self.superfluous_keys
if len(superfluous) == 0:
for key in data.keys():
if key not in self.required_keys:
superfluous.add(key)
for key in superfluous:
del data[key]
return data
class SerializeBox(DataProcess):
box_key = State(default='charboxes')
format = State(default='NP2')
def process(self, data):
data[self.box_key] = data['lines'].quads
return data
class MultiStepLR(Configurable):
lr = State()
milestones = State(default=[]) # milestones must be sorted
gamma = State(default=0.1)
def __init__(self, cmd={}, **kwargs):
self.load_all(**kwargs)
self.lr = cmd.get('lr', self.lr)
def get_learning_rate(self, epoch, step):
return self.lr * self.gamma**bisect_right(self.milestones, epoch)
class DecayLearningRate(Configurable):
lr = State(default=0.007)
epochs = State(default=1200)
factor = State(default=0.9)
def __init__(self, **kwargs):
self.load_all(**kwargs)
def get_learning_rate(self, epoch, step=None):
rate = np.power(1.0 - epoch / float(self.epochs + 1), self.factor)
return rate * self.lr
class WarmupLR(Configurable):
steps = State(default=4000)
warmup_lr = State(default=1e-5)
origin_lr = State()
def __init__(self, cmd={}, **kwargs):
self.load_all(**kwargs)
def get_learning_rate(self, epoch, step):
if epoch == 0 and step < self.steps:
return self.warmup_lr
return self.origin_lr.get_learning_rate(epoch, step)
class TrainSettings(Configurable):
data_loader = State()
model_saver = State()
checkpoint = State()
scheduler = State()
epochs = State(default=10)
def __init__(self, **kwargs):
kwargs['cmd'].update(is_train=True)
self.load_all(**kwargs)
if 'epochs' in kwargs['cmd']:
self.epochs = kwargs['cmd']['epochs']
class ValidationSettings(Configurable):
data_loaders = State()
visualize = State()
interval = State(default=100)
exempt = State(default=-1)
def __init__(self, **kwargs):
kwargs['cmd'].update(is_train=False)
self.load_all(**kwargs)
cmd = kwargs['cmd']
self.visualize = cmd['visualize']
class PiecewiseConstantLearningRate(Configurable):
boundaries = State(default=[10000, 20000])
values = State(default=[0.001, 0.0001, 0.00001])
def __init__(self, **kwargs):
self.load_all(**kwargs)
def get_learning_rate(self, epoch, step):
for boundary, value in zip(self.boundaries, self.values[:-1]):
if step < boundary:
return value
return self.values[-1]
class Structure(Configurable):
builder = State()
representer = State()
measurer = State()
visualizer = State()
def __init__(self, **kwargs):
self.load_all(**kwargs)
@property
def model_name(self):
return self.builder.model_name
class MakeCenterMap(DataProcess):
max_size = State(default=32)
shape = State(default=(64, 256))
sigma_ratio = State(default=16)
function_name = State(default='sample_gaussian')
points_key = 'points'
correlation = 0 # The formulation of guassian is simplified when correlation is 0
def process(self, data):
assert self.points_key in data, '%s in data is required' % self.points_key
points = data['points'] * self.shape[::-1] # N, 2
assert points.shape[0] >= self.max_size
func = getattr(self, self.function_name)
data['charmaps'] = func(points, *self.shape)
return data
def gaussian(self, points, height, width):
index_x, index_y = np.meshgrid(np.linspace(0, width, width),
np.linspace(0, height, height))
index_x = np.repeat(index_x[np.newaxis], points.shape[0], axis=0)
index_y = np.repeat(index_y[np.newaxis], points.shape[0], axis=0)
mu_x = points[:, 0][:, np.newaxis, np.newaxis]
mu_y = points[:, 1][:, np.newaxis, np.newaxis]
mask_is_zero = ((mu_x == 0) + (mu_y == 0)) == 0
result = np.reciprocal(2 * np.pi * width / self.sigma_ratio * height / self.sigma_ratio)\
* np.exp(- 0.5 * (np.square((index_x - mu_x) / width * self.sigma_ratio) +
np.square((index_y - mu_y) / height * self.sigma_ratio)))
result = result / \
np.maximum(result.max(axis=1, keepdims=True).max(
axis=2, keepdims=True), np.finfo(np.float32).eps)
result = result * mask_is_zero
return result.astype(np.float32)
def sample_gaussian(self, points, height, width):
points = (points + 0.5).astype(np.int32)
canvas = np.zeros((self.max_size, height, width), dtype=np.float32)
for index in range(canvas.shape[0]):
point = points[index]
canvas[index, point[1], point[0]] = 1.
if point.sum() > 0:
fi.gaussian_filter(canvas[index], (height // self.sigma_ratio,
width // self.sigma_ratio),
output=canvas[index], mode='mirror')
canvas[index] = canvas[index] / canvas[index].max()
x_range = min(point[0], width - point[0])
canvas[index, :, :point[0] - x_range] = 0
canvas[index, :, point[0] + x_range:] = 0
y_range = min(point[1], width - point[1])
canvas[index, :point[1] - y_range, :] = 0
canvas[index, point[1] + y_range:, :] = 0
return canvas
class MakeCenterPoints(DataProcess):
box_key = State(default='charboxes')
size = State(default=32)
def process(self, data):
shape = data['image'].shape[:2]
points = np.zeros((self.size, 2), dtype=np.float32)
boxes = np.array(data[self.box_key])[:self.size]
size = boxes.shape[0]
points[:size] = boxes.mean(axis=1)
data['points'] = (points / shape[::-1]).astype(np.float32)
return data
class ResizeImage(_ResizeImage, DataProcess):
mode = State(default='keep_ratio')
image_size = State(default=[1152, 2048]) # height, width
key = State(default='image')
def __init__(self, cmd={}, mode=None, **kwargs):
self.load_all(**kwargs)
if mode is not None:
self.mode = mode
if 'resize_mode' in cmd:
self.mode = cmd['resize_mode']
assert self.mode in self.MODES
def process(self, data):
data[self.key] = self.resize_or_pad(data[self.key])
return data
class ConstantLearningRate(Configurable):
lr = State(default=0.0001)
def __init__(self, **kwargs):
self.load_all(**kwargs)
def get_learning_rate(self, epoch, step):
return self.lr
class OptimizerScheduler(Configurable):
optimizer = State()
optimizer_args = State(default={})
learning_rate = State(autoload=False)
def __init__(self, cmd={}, **kwargs):
self.load_all(**kwargs)
self.load('learning_rate', cmd=cmd, **kwargs)
if 'lr' in cmd:
self.optimizer_args['lr'] = cmd['lr']
def create_optimizer(self, parameters):
optimizer = getattr(torch.optim, self.optimizer)(parameters,
**self.optimizer_args)
if hasattr(self.learning_rate, 'prepare'):
self.learning_rate.prepare(optimizer)
return optimizer
class Builder(Configurable):
model = State()
model_args = State()
def __init__(self, cmd={}, **kwargs):
self.load_all(**kwargs)
if 'backbone' in cmd:
self.model_args['backbone'] = cmd['backbone']
@property
def model_name(self):
return self.model + '-' + getattr(structure_model, self.model).model_name(self.model_args)
def build(self, device, distributed=False, local_rank: int = 0):
Model = getattr(structure_model,self.model)
model = Model(self.model_args, device,
distributed=distributed, local_rank=local_rank)
return model
class SliceDataset(TorchDataset, Configurable):
dataset = State()
start = State()
end = State()
def __init__(self, **kwargs):
self.load_all(**kwargs)
if self.start is None:
self.start = 0
if self.end is None:
self.end = len(self.dataset)
def __getitem__(self, idx):
return self.dataset[self.start + idx]
def __len__(self):
return self.end - self.start
class PriorityLearningRate(Configurable):
learning_rates = State()
def __init__(self, **kwargs):
self.load_all(**kwargs)
def get_learning_rate(self, epoch, step):
for learning_rate in self.learning_rates:
lr = learning_rate.get_learning_rate(epoch, step)
if lr is not None:
return lr
return None
class Experiment(Configurable):
structure = State(autoload=False)
train = State()
validation = State(autoload=False)
evaluation = State(autoload=False)
logger = State(autoload=True)
def __init__(self, **kwargs):
self.load('structure', **kwargs)
cmd = kwargs.get('cmd', {})
if 'name' not in cmd:
cmd['name'] = self.structure.model_name
self.load_all(**kwargs)
self.distributed = cmd.get('distributed', False)
self.local_rank = cmd.get('local_rank', 0)
if cmd.get('validate', False):
self.load('validation', **kwargs)
else:
self.validation = None
class FileMonitorLearningRate(Configurable):
file_path = State()
def __init__(self, **kwargs):
self.load_all(**kwargs)
self.monitor = SignalMonitor(self.file_path)
def get_learning_rate(self, epoch, step):
signal = self.monitor.get_signal()
if signal is not None:
return float(signal)
return None
class InfiniteDataLoader(Configurable, torch.utils.data.DataLoader):
dataset = State()
batch_size = State(default=256)
num_workers = State(default=10)
limit_size = State(default=2**31)
def __init__(self, **kwargs):
self.load_all(**kwargs)
cmd = kwargs['cmd']
if 'batch_size' in cmd:
self.batch_size = cmd['batch_size']
sampler = InfiniteOrderedSampler(self.dataset, self.limit_size)
torch.utils.data.DataLoader.__init__(
self,
self.dataset,
batch_size=self.batch_size,
num_workers=self.num_workers,
sampler=sampler,
worker_init_fn=default_worker_init_fn,
)
class ModelSaver(Configurable):
dir_path = State()
save_interval = State(default=1000)
signal_path = State()
def __init__(self, **kwargs):
self.load_all(**kwargs)
# BUG: signal path should not be global
self.monitor = SignalMonitor(self.signal_path)
def maybe_save_model(self, model, epoch, step, logger):
if step % self.save_interval == 0 or self.monitor.get_signal(
) is not None:
self.save_model(model, epoch, step)
logger.report_time('Saving ')
logger.iter(step)
def save_model(self, model, epoch=None, step=None):
if isinstance(model, dict):
for name, net in model.items():
checkpoint_name = self.make_checkpoint_name(name, epoch, step)
self.save_checkpoint(net, checkpoint_name)
else:
checkpoint_name = self.make_checkpoint_name('model', epoch, step)
self.save_checkpoint(model, checkpoint_name)
def save_checkpoint(self, net, name):
os.makedirs(self.dir_path, exist_ok=True)
torch.save(net.state_dict(), os.path.join(self.dir_path, name))
def make_checkpoint_name(self, name, epoch=None, step=None):
if epoch is None or step is None:
c_name = name + '_latest'
else:
c_name = '{}_epoch_{}_minibatch_{}'.format(name, epoch, step)
return c_name
class BuitlinLearningRate(Configurable):
lr = State(default=0.001)
klass = State(default='StepLR')
args = State(default=[])
kwargs = State(default={})
def __init__(self, cmd={}, **kwargs):
self.load_all(**kwargs)
self.lr = cmd.get('lr', None) or self.lr
self.scheduler = None
def prepare(self, optimizer):
self.scheduler = getattr(lr_scheduler,
self.klass)(optimizer, *self.args,
**self.kwargs)
def get_learning_rate(self, epoch, step=None):
if self.scheduler is None:
raise 'learning rate not ready(prepared with optimizer) '
self.scheduler.last_epoch = epoch
# return value of gt_lr is a list,
# where each element is the corresponding learning rate for a
# paramater group.
return self.scheduler.get_lr()[0]
class DataIdMetaLoader(MetaLoader):
return_dict = State(default=False)
scan_meta = False
def __init__(self, return_dict=None, cmd={}, **kwargs):
super().__init__(cmd=cmd, **kwargs)
if return_dict is not None:
self.return_dict = return_dict
def parse_meta(self, data_id):
return dict(data_id=data_id)
def post_prosess(self, meta):
if self.return_dict:
return meta
return meta['data_id']
class UnifyRect(SerializeBox):
max_size = State(default=64)
def process(self, data):
h, w = data['image'].shape[:2]
boxes = np.zeros((self.max_size, 4), dtype=np.float32)
mask_has_box = np.zeros(self.max_size, dtype=np.float32)
data = super().process(data)
quad = data[self.box_key]
assert quad.shape[0] <= self.max_size
boxes[:quad.shape[0]] = quad.rectify() / np.array([w, h, w, h
]).reshape(1, 4)
mask_has_box[:quad.shape[0]] = 1.
data['boxes'] = boxes
data['mask_has_box'] = mask_has_box
return data
class RandomCropAug(Configurable):
size = State(default=640)
def __init__(self, size=640, *args, **kwargs):
self.size = size or self.size
self.augment = RandomCrop(size)
def __call__(self, data):
'''
This augmenter is supposed to following the process of `MakeICDARData`,
in which labels are mapped to this specific format:
(image, polygons: (n, 4, 2), tags: [Boolean], ...)
'''
image, boxes, ignore_tags = data[:3]
image, boxes, ignore_tags = self.augment(image, boxes, ignore_tags)
return (image, boxes, ignore_tags, *data[3:])