def train(self, tmp_dir):
"""Train a model."""
self.print_options()
# Sync output of previous training run from cloud.
train_uri = self.backend_opts.train_uri
train_dir = get_local_path(train_uri, tmp_dir)
make_dir(train_dir)
sync_from_dir(train_uri, train_dir)
'''
Get zip file for each group, and unzip them into chip_dir in a
way that works well with FastAI.
The resulting directory structure would be:
<chip_dir>/
train/
training-<uuid1>/
<class1>/
...
<class2>/
...
...
training-<uuid2>/
<class1>/
...
<class2>/
...
...
...
val/
validation-<uuid1>/
<class1>/
...
<class2>/
...
...
validation-<uuid2>/
<class1>/
...
<class2>/
...
...
...
'''
chip_dir = join(tmp_dir, 'chips/')
make_dir(chip_dir)
for zip_uri in list_paths(self.backend_opts.chip_uri, 'zip'):
zip_name = Path(zip_uri).name
if zip_name.startswith('train'):
extract_dir = chip_dir + 'train/'
elif zip_name.startswith('val'):
extract_dir = chip_dir + 'val/'
else:
continue
zip_path = download_if_needed(zip_uri, tmp_dir)
with zipfile.ZipFile(zip_path, 'r') as zipf:
zipf.extractall(extract_dir)
# Setup data loader.
def get_label_path(im_path):
return Path(str(im_path.parent)[:-4] + '-labels') / im_path.name
size = self.task_config.chip_size
class_map = self.task_config.class_map
classes = class_map.get_class_names()
num_workers = 0 if self.train_opts.debug else 4
tfms = get_transforms(flip_vert=self.train_opts.flip_vert)
def get_data(train_sampler=None):
data = (ImageList.from_folder(chip_dir).split_by_folder(
train='train', valid='val').label_from_folder().transform(
tfms, size=size).databunch(
bs=self.train_opts.batch_sz,
num_workers=num_workers,
))
return data
data = get_data()
if self.train_opts.debug:
make_debug_chips(data, class_map, tmp_dir, train_uri)
# Setup learner.
ignore_idx = -1
metrics = [
Precision(average='weighted', clas_idx=1, ignore_idx=ignore_idx),
Recall(average='weighted', clas_idx=1, ignore_idx=ignore_idx),
FBeta(average='weighted',
clas_idx=1,
beta=1,
ignore_idx=ignore_idx)
]
model_arch = getattr(models, self.train_opts.model_arch)
learn = cnn_learner(data,
model_arch,
metrics=metrics,
wd=self.train_opts.weight_decay,
path=train_dir)
learn.unfreeze()
if self.train_opts.fp16 and torch.cuda.is_available():
# This loss_scale works for Resnet 34 and 50. You might need to adjust this
# for other models.
learn = learn.to_fp16(loss_scale=256)
# Setup callbacks and train model.
model_path = get_local_path(self.backend_opts.model_uri, tmp_dir)
pretrained_uri = self.backend_opts.pretrained_uri
if pretrained_uri:
print('Loading weights from pretrained_uri: {}'.format(
pretrained_uri))
pretrained_path = download_if_needed(pretrained_uri, tmp_dir)
learn.model.load_state_dict(torch.load(
pretrained_path, map_location=learn.data.device),
strict=False)
# Save every epoch so that resume functionality provided by
# TrackEpochCallback will work.
callbacks = [
TrackEpochCallback(learn),
MySaveModelCallback(learn, every='epoch'),
MyCSVLogger(learn, filename='log'),
ExportCallback(learn, model_path, monitor='f_beta'),
SyncCallback(train_dir, self.backend_opts.train_uri,
self.train_opts.sync_interval)
]
lr = self.train_opts.lr
num_epochs = self.train_opts.num_epochs
if self.train_opts.one_cycle:
if lr is None:
learn.lr_find()
learn.recorder.plot(suggestion=True, return_fig=True)
lr = learn.recorder.min_grad_lr
print('lr_find() found lr: {}'.format(lr))
learn.fit_one_cycle(num_epochs, lr, callbacks=callbacks)
else:
learn.fit(num_epochs, lr, callbacks=callbacks)
# Since model is exported every epoch, we need some other way to
# show that training is finished.
str_to_file('done!', self.backend_opts.train_done_uri)
# Sync output to cloud.
sync_to_dir(train_dir, self.backend_opts.train_uri)
def train(self, tmp_dir):
"""Train a model."""
self.print_options()
# Sync output of previous training run from cloud.
train_uri = self.backend_opts.train_uri
train_dir = get_local_path(train_uri, tmp_dir)
make_dir(train_dir)
sync_from_dir(train_uri, train_dir)
# Get zip file for each group, and unzip them into chip_dir.
chip_dir = join(tmp_dir, 'chips')
make_dir(chip_dir)
for zip_uri in list_paths(self.backend_opts.chip_uri, 'zip'):
zip_path = download_if_needed(zip_uri, tmp_dir)
with zipfile.ZipFile(zip_path, 'r') as zipf:
zipf.extractall(chip_dir)
# Setup data loader.
train_images = []
train_lbl_bbox = []
for annotation_path in glob.glob(join(chip_dir, 'train/*.json')):
images, lbl_bbox = get_annotations(annotation_path)
train_images += images
train_lbl_bbox += lbl_bbox
val_images = []
val_lbl_bbox = []
for annotation_path in glob.glob(join(chip_dir, 'valid/*.json')):
images, lbl_bbox = get_annotations(annotation_path)
val_images += images
val_lbl_bbox += lbl_bbox
images = train_images + val_images
lbl_bbox = train_lbl_bbox + val_lbl_bbox
img2bbox = dict(zip(images, lbl_bbox))
get_y_func = lambda o: img2bbox[o.name]
num_workers = 0 if self.train_opts.debug else 4
data = ObjectItemList.from_folder(chip_dir)
data = data.split_by_folder()
data = data.label_from_func(get_y_func)
data = data.transform(
get_transforms(), size=self.task_config.chip_size, tfm_y=True)
data = data.databunch(
bs=self.train_opts.batch_sz, collate_fn=bb_pad_collate,
num_workers=num_workers)
print(data)
if self.train_opts.debug:
make_debug_chips(
data, self.task_config.class_map, tmp_dir, train_uri)
# Setup callbacks and train model.
ratios = [1/2, 1, 2]
scales = [1, 2**(-1/3), 2**(-2/3)]
model_arch = getattr(models, self.train_opts.model_arch)
encoder = create_body(model_arch, cut=-2)
model = RetinaNet(encoder, data.c, final_bias=-4)
crit = RetinaNetFocalLoss(scales=scales, ratios=ratios)
learn = Learner(data, model, loss_func=crit, path=train_dir)
learn = learn.split(retina_net_split)
model_path = get_local_path(self.backend_opts.model_uri, tmp_dir)
pretrained_uri = self.backend_opts.pretrained_uri
if pretrained_uri:
print('Loading weights from pretrained_uri: {}'.format(
pretrained_uri))
pretrained_path = download_if_needed(pretrained_uri, tmp_dir)
learn.load(pretrained_path[:-4])
callbacks = [
TrackEpochCallback(learn),
SaveModelCallback(learn, every='epoch'),
MyCSVLogger(learn, filename='log'),
ExportCallback(learn, model_path),
SyncCallback(train_dir, self.backend_opts.train_uri,
self.train_opts.sync_interval)
]
learn.unfreeze()
learn.fit(self.train_opts.num_epochs, self.train_opts.lr,
callbacks=callbacks)
# Since model is exported every epoch, we need some other way to
# show that training is finished.
str_to_file('done!', self.backend_opts.train_done_uri)
# Sync output to cloud.
sync_to_dir(train_dir, self.backend_opts.train_uri)
def train(self, tmp_dir):
"""Train a model."""
self.print_options()
# Sync output of previous training run from cloud.
train_uri = self.backend_opts.train_uri
train_dir = get_local_path(train_uri, tmp_dir)
make_dir(train_dir)
sync_from_dir(train_uri, train_dir)
# Get zip file for each group, and unzip them into chip_dir.
chip_dir = join(tmp_dir, 'chips')
make_dir(chip_dir)
for zip_uri in list_paths(self.backend_opts.chip_uri, 'zip'):
zip_path = download_if_needed(zip_uri, tmp_dir)
with zipfile.ZipFile(zip_path, 'r') as zipf:
zipf.extractall(chip_dir)
# Setup data loader.
def get_label_path(im_path):
return Path(str(im_path.parent)[:-4] + '-labels') / im_path.name
size = self.task_config.chip_size
class_map = self.task_config.class_map
classes = class_map.get_class_names()
if 0 not in class_map.get_keys():
classes = ['nodata'] + classes
num_workers = 0 if self.train_opts.debug else 4
train_img_dir = self.subset_training_data(chip_dir)
def get_data(train_sampler=None):
data = (SegmentationItemList.from_folder(chip_dir).split_by_folder(
train=train_img_dir, valid='val-img').label_from_func(
get_label_path, classes=classes).transform(
get_transforms(flip_vert=self.train_opts.flip_vert),
size=size,
tfm_y=True).databunch(bs=self.train_opts.batch_sz,
num_workers=num_workers,
train_sampler=train_sampler))
return data
data = get_data()
oversample = self.train_opts.oversample
if oversample:
sampler = get_weighted_sampler(data.train_ds,
oversample['rare_class_ids'],
oversample['rare_target_prop'])
data = get_data(train_sampler=sampler)
if self.train_opts.debug:
make_debug_chips(data, class_map, tmp_dir, train_uri)
# Setup learner.
ignore_idx = 0
metrics = [
Precision(average='weighted', clas_idx=1, ignore_idx=ignore_idx),
Recall(average='weighted', clas_idx=1, ignore_idx=ignore_idx),
FBeta(average='weighted',
clas_idx=1,
beta=1,
ignore_idx=ignore_idx)
]
model_arch = getattr(models, self.train_opts.model_arch)
learn = unet_learner(data,
model_arch,
metrics=metrics,
wd=self.train_opts.weight_decay,
bottle=True,
path=train_dir)
learn.unfreeze()
if self.train_opts.fp16 and torch.cuda.is_available():
# This loss_scale works for Resnet 34 and 50. You might need to adjust this
# for other models.
learn = learn.to_fp16(loss_scale=256)
# Setup callbacks and train model.
model_path = get_local_path(self.backend_opts.model_uri, tmp_dir)
pretrained_uri = self.backend_opts.pretrained_uri
if pretrained_uri:
print('Loading weights from pretrained_uri: {}'.format(
pretrained_uri))
pretrained_path = download_if_needed(pretrained_uri, tmp_dir)
learn.model.load_state_dict(torch.load(
pretrained_path, map_location=learn.data.device),
strict=False)
# Save every epoch so that resume functionality provided by
# TrackEpochCallback will work.
callbacks = [
TrackEpochCallback(learn),
MySaveModelCallback(learn, every='epoch'),
MyCSVLogger(learn, filename='log'),
ExportCallback(learn, model_path, monitor='f_beta'),
SyncCallback(train_dir, self.backend_opts.train_uri,
self.train_opts.sync_interval)
]
lr = self.train_opts.lr
num_epochs = self.train_opts.num_epochs
if self.train_opts.one_cycle:
if lr is None:
learn.lr_find()
learn.recorder.plot(suggestion=True, return_fig=True)
lr = learn.recorder.min_grad_lr
print('lr_find() found lr: {}'.format(lr))
learn.fit_one_cycle(num_epochs, lr, callbacks=callbacks)
else:
learn.fit(num_epochs, lr, callbacks=callbacks)
# Since model is exported every epoch, we need some other way to
# show that training is finished.
str_to_file('done!', self.backend_opts.train_done_uri)
# Sync output to cloud.
sync_to_dir(train_dir, self.backend_opts.train_uri)
def train(self, tmp_dir):
"""Train a model.
This downloads any previous output saved to the train_uri,
starts training (or resumes from a checkpoint), periodically
syncs contents of train_dir to train_uri and after training finishes.
Args:
tmp_dir: (str) path to temp directory
"""
self.log_options()
# Sync output of previous training run from cloud.
train_uri = self.backend_opts.train_uri
train_dir = get_local_path(train_uri, tmp_dir)
make_dir(train_dir)
sync_from_dir(train_uri, train_dir)
# Get zip file for each group, and unzip them into chip_dir.
chip_dir = join(tmp_dir, 'chips')
make_dir(chip_dir)
for zip_uri in list_paths(self.backend_opts.chip_uri, 'zip'):
zip_path = download_if_needed(zip_uri, tmp_dir)
with zipfile.ZipFile(zip_path, 'r') as zipf:
zipf.extractall(chip_dir)
# Setup data loader.
def get_label_path(im_path):
return Path(str(im_path.parent)[:-4] + '-labels') / im_path.name
size = self.task_config.chip_size
class_map = self.task_config.class_map
classes = class_map.get_class_names()
if 0 not in class_map.get_keys():
classes = ['nodata'] + classes
num_workers = 0 if self.train_opts.debug else 4
data = (SegmentationItemList.from_folder(chip_dir)
.split_by_folder(train='train-img', valid='val-img'))
train_count = None
if self.train_opts.train_count is not None:
train_count = min(len(data.train), self.train_opts.train_count)
elif self.train_opts.train_prop != 1.0:
train_count = int(round(self.train_opts.train_prop * len(data.train)))
train_items = data.train.items
if train_count is not None:
train_inds = np.random.permutation(np.arange(len(data.train)))[0:train_count]
train_items = train_items[train_inds]
items = np.concatenate([train_items, data.valid.items])
data = (SegmentationItemList(items, chip_dir)
.split_by_folder(train='train-img', valid='val-img')
.label_from_func(get_label_path, classes=classes)
.transform(get_transforms(flip_vert=self.train_opts.flip_vert),
size=size, tfm_y=True)
.databunch(bs=self.train_opts.batch_sz,
num_workers=num_workers))
print(data)
# Setup learner.
ignore_idx = 0
metrics = [
Precision(average='weighted', clas_idx=1, ignore_idx=ignore_idx),
Recall(average='weighted', clas_idx=1, ignore_idx=ignore_idx),
FBeta(average='weighted', clas_idx=1, beta=1, ignore_idx=ignore_idx)]
model_arch = getattr(models, self.train_opts.model_arch)
learn = unet_learner(
data, model_arch, metrics=metrics, wd=self.train_opts.weight_decay,
bottle=True, path=train_dir)
learn.unfreeze()
if self.train_opts.mixed_prec and torch.cuda.is_available():
# This loss_scale works for Resnet 34 and 50. You might need to adjust this
# for other models.
learn = learn.to_fp16(loss_scale=256)
# Setup callbacks and train model.
model_path = get_local_path(self.backend_opts.model_uri, tmp_dir)
pretrained_uri = self.backend_opts.pretrained_uri
if pretrained_uri:
print('Loading weights from pretrained_uri: {}'.format(
pretrained_uri))
pretrained_path = download_if_needed(pretrained_uri, tmp_dir)
learn.model.load_state_dict(
torch.load(pretrained_path, map_location=learn.data.device),
strict=False)
# Save every epoch so that resume functionality provided by
# TrackEpochCallback will work.
callbacks = [
TrackEpochCallback(learn),
MySaveModelCallback(learn, every='epoch'),
MyCSVLogger(learn, filename='log'),
ExportCallback(learn, model_path, monitor='f_beta'),
SyncCallback(train_dir, self.backend_opts.train_uri,
self.train_opts.sync_interval)
]
oversample = self.train_opts.oversample
if oversample:
weights = get_oversampling_weights(
data.train_ds, oversample['rare_class_ids'],
oversample['rare_target_prop'])
oversample_callback = OverSamplingCallback(learn, weights=weights)
callbacks.append(oversample_callback)
if self.train_opts.debug:
if oversample:
oversample_callback.on_train_begin()
make_debug_chips(data, class_map, tmp_dir, train_uri)
if self.train_opts.log_tensorboard:
callbacks.append(TensorboardLogger(learn, 'run'))
if self.train_opts.run_tensorboard:
log.info('Starting tensorboard process')
log_dir = join(train_dir, 'logs', 'run')
tensorboard_process = Popen(
['tensorboard', '--logdir={}'.format(log_dir)])
terminate_at_exit(tensorboard_process)
lr = self.train_opts.lr
num_epochs = self.train_opts.num_epochs
if self.train_opts.one_cycle:
if lr is None:
learn.lr_find()
learn.recorder.plot(suggestion=True, return_fig=True)
lr = learn.recorder.min_grad_lr
print('lr_find() found lr: {}'.format(lr))
learn.fit_one_cycle(num_epochs, lr, callbacks=callbacks)
else:
learn.fit(num_epochs, lr, callbacks=callbacks)
if self.train_opts.run_tensorboard:
tensorboard_process.terminate()
# Since model is exported every epoch, we need some other way to
# show that training is finished.
str_to_file('done!', self.backend_opts.train_done_uri)
# Sync output to cloud.
sync_to_dir(train_dir, self.backend_opts.train_uri)
def train(config_path, opts):
tmp_dir_obj = tempfile.TemporaryDirectory()
tmp_dir = tmp_dir_obj.name
cfg = load_config(config_path, opts)
print(cfg)
# Setup data
databunch, full_databunch = build_databunch(cfg, tmp_dir)
output_dir = setup_output_dir(cfg, tmp_dir)
print(full_databunch)
plotter = build_plotter(cfg)
if not cfg.lr_find_mode and not cfg.predict_mode:
plotter.plot_data(databunch, output_dir)
# Setup model
num_labels = databunch.c
model = build_model(cfg, num_labels)
metrics = [CocoMetric(num_labels)]
learn = Learner(databunch, model, path=output_dir, metrics=metrics)
fastai.basic_train.loss_batch = loss_batch
best_model_path = join(output_dir, 'best_model.pth')
last_model_path = join(output_dir, 'last_model.pth')
# Train model
callbacks = [
MyCSVLogger(learn, filename='log'),
SubLossMetric(learn, model.subloss_names)
]
if cfg.output_uri.startswith('s3://'):
callbacks.append(
SyncCallback(output_dir, cfg.output_uri, cfg.solver.sync_interval))
if cfg.model.init_weights:
device = next(model.parameters()).device
model.load_state_dict(
torch.load(cfg.model.init_weights, map_location=device))
if not cfg.predict_mode:
if cfg.overfit_mode:
learn.fit_one_cycle(cfg.solver.num_epochs, cfg.solver.lr, callbacks=callbacks)
torch.save(learn.model.state_dict(), best_model_path)
learn.model.eval()
print('Validating on training set...')
learn.validate(full_databunch.train_dl, metrics=metrics)
else:
tb_logger = TensorboardLogger(learn, 'run')
tb_logger.set_extra_args(
model.subloss_names, cfg.overfit_mode)
extra_callbacks = [
MySaveModelCallback(
learn, best_model_path, monitor='coco_metric', every='improvement'),
MySaveModelCallback(learn, last_model_path, every='epoch'),
TrackEpochCallback(learn),
]
callbacks.extend(extra_callbacks)
if cfg.lr_find_mode:
learn.lr_find()
learn.recorder.plot(suggestion=True, return_fig=True)
lr = learn.recorder.min_grad_lr
print('lr_find() found lr: {}'.format(lr))
exit()
learn.fit_one_cycle(cfg.solver.num_epochs, cfg.solver.lr, callbacks=callbacks)
print('Validating on full validation set...')
learn.validate(full_databunch.valid_dl, metrics=metrics)
else:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.load_state_dict(
torch.load(join(output_dir, 'best_model.pth'), map_location=device))
model.eval()
plot_dataset = databunch.train_ds
print('Plotting predictions...')
plot_dataset = databunch.train_ds if cfg.overfit_mode else databunch.valid_ds
plotter.make_debug_plots(plot_dataset, model, databunch.classes, output_dir)
if cfg.output_uri.startswith('s3://'):
sync_to_dir(output_dir, cfg.output_uri)