def collect_experiment(key, root_uri, output_dir, get_pred_package=False):
print('\nCollecting experiment {}...\n'.format(key))
if root_uri.startswith('s3://'):
predict_package_uris = list_paths(join(root_uri, key, 'bundle'), ext='predict_package.zip')
eval_json_uris = list_paths(join(root_uri, key, 'eval'), ext='eval.json')
else:
predict_package_uris = glob.glob(join(root_uri, key, 'bundle', '*', 'predict_package.zip'))
eval_json_uris = glob.glob(join(root_uri, key, 'eval', '*', 'eval.json'))
if len(predict_package_uris) > 1 or len(eval_json_uris) > 1:
print('Cannot collect from key with multiple experiments!!!')
return
if len(predict_package_uris) == 0 or len(eval_json_uris) == 0:
print('Missing output!!!')
return
predict_package_uri = predict_package_uris[0]
eval_json_uri = eval_json_uris[0]
make_dir(join(output_dir, key))
if get_pred_package:
download_or_copy(predict_package_uri, join(output_dir, key))
download_or_copy(eval_json_uri, join(output_dir, key))
eval_json = file_to_json(join(output_dir, key, 'eval.json'))
pprint.pprint(eval_json['overall'], indent=4)
def __enter__(self):
self.tmp_dir_obj = tempfile.TemporaryDirectory(dir=self.tmp_dir)
self.sample_dir = join(self.tmp_dir_obj.name, 'samples')
make_dir(self.sample_dir)
self.sample_ind = 0
return self
def setup_tensorboard(self):
"""Setup for logging stats to TB."""
self.tb_writer = None
if self.cfg.log_tensorboard:
self.tb_log_dir = join(self.output_dir, 'tb-logs')
make_dir(self.tb_log_dir)
self.tb_writer = SummaryWriter(log_dir=self.tb_log_dir)
def plot_batch(self, x: Tensor, y, output_path: str, z=None):
"""Plot a whole batch in a grid using plot_xyz.
Args:
x: batch of images
y: ground truth labels
output_path: local path where to save plot image
z: optional predicted labels
"""
batch_sz = x.shape[0]
ncols = nrows = math.ceil(math.sqrt(batch_sz))
fig = plt.figure(constrained_layout=True,
figsize=(3 * ncols, 3 * nrows))
grid = gridspec.GridSpec(ncols=ncols, nrows=nrows, figure=fig)
for i in range(batch_sz):
ax = fig.add_subplot(grid[i])
if z is None:
self.plot_xyz(ax, x[i], y[i])
else:
self.plot_xyz(ax, x[i], y[i], z=z[i])
make_dir(output_path, use_dirname=True)
plt.savefig(output_path)
plt.close()
def save_model_bundle(self):
model_bundle_dir = join(self.tmp_dir, 'model-bundle')
make_dir(model_bundle_dir)
shutil.copyfile(self.last_model_path,
join(model_bundle_dir, 'model.pth'))
shutil.copyfile(self.config_path, join(model_bundle_dir,
'config.json'))
zipdir(model_bundle_dir, self.model_bundle_path)
def test_file_exists_local_true(self):
path = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum.txt')
directory = os.path.dirname(path)
make_dir(directory, check_empty=False)
str_to_file(self.lorem, path)
self.assertTrue(file_exists(path))
def test_check_empty(self):
path = os.path.join(self.tmp_dir.name, 'hello', 'hello.txt')
dir = os.path.dirname(path)
str_to_file('hello', path)
make_dir(dir, check_empty=False)
with self.assertRaises(Exception):
make_dir(dir, check_empty=True)
def test_sync_from_dir_noop_local(self):
path = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum.txt')
src = os.path.join(self.tmp_dir.name, 'lorem')
make_dir(src, check_empty=False)
fs = FileSystem.get_file_system(src, 'r')
fs.write_bytes(path, bytes([0x00, 0x01]))
sync_from_dir(src, src, delete=True)
self.assertEqual(len(list_paths(src)), 1)
def test_copy_to_http(self):
path = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum.txt')
dst = 'http://localhost/'
directory = os.path.dirname(path)
make_dir(directory, check_empty=False)
str_to_file(self.lorem, path)
self.assertRaises(NotWritableError, lambda: upload_or_copy(path, dst))
os.remove(path)
def test_force_empty(self):
path = os.path.join(self.tmp_dir.name, 'hello', 'hello.txt')
dir = os.path.dirname(path)
str_to_file('hello', path)
make_dir(dir, force_empty=False)
self.assertTrue(os.path.isfile(path))
make_dir(dir, force_empty=True)
is_empty = len(os.listdir(dir)) == 0
self.assertTrue(is_empty)
def test_last_modified(self):
path = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum1.txt')
directory = os.path.dirname(path)
make_dir(directory, check_empty=False)
fs = FileSystem.get_file_system(path, 'r')
str_to_file(self.lorem, path)
stamp = fs.last_modified(path)
self.assertTrue(isinstance(stamp, datetime.datetime))
def test_file_exists_s3_true(self):
path = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum.txt')
directory = os.path.dirname(path)
make_dir(directory, check_empty=False)
str_to_file(self.lorem, path)
s3_path = 's3://{}/lorem.txt'.format(self.bucket_name)
upload_or_copy(path, s3_path)
self.assertTrue(file_exists(s3_path))
def test_copy_to_local(self):
path1 = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum.txt')
path2 = os.path.join(self.tmp_dir.name, 'yyy', 'ipsum.txt')
dir1 = os.path.dirname(path1)
dir2 = os.path.dirname(path2)
make_dir(dir1, check_empty=False)
make_dir(dir2, check_empty=False)
str_to_file(self.lorem, path1)
upload_or_copy(path1, path2)
self.assertEqual(len(list_paths(dir2)), 1)
def test_sync_to_dir_local(self):
path = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum.txt')
src = os.path.dirname(path)
dst = os.path.join(self.tmp_dir.name, 'xxx')
make_dir(src, check_empty=False)
make_dir(dst, check_empty=False)
fs = FileSystem.get_file_system(path, 'r')
fs.write_bytes(path, bytes([0x00, 0x01]))
sync_to_dir(src, dst, delete=True)
self.assertEqual(len(list_paths(dst)), 1)
def test_list_paths_s3(self):
path = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum.txt')
s3_path = 's3://{}/xxx/lorem.txt'.format(self.bucket_name)
s3_directory = 's3://{}/xxx/'.format(self.bucket_name)
directory = os.path.dirname(path)
make_dir(directory, check_empty=False)
str_to_file(self.lorem, path)
upload_or_copy(path, s3_path)
list_paths(s3_directory)
self.assertEqual(len(list_paths(s3_directory)), 1)
def test_bytes_local(self):
path = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum.txt')
directory = os.path.dirname(path)
make_dir(directory, check_empty=False)
expected = bytes([0x00, 0x01, 0x02])
fs = FileSystem.get_file_system(path, 'r')
fs.write_bytes(path, expected)
actual = fs.read_bytes(path)
self.assertEqual(actual, expected)
def test_last_modified_s3(self):
path = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum1.txt')
s3_path = 's3://{}/lorem1.txt'.format(self.bucket_name)
directory = os.path.dirname(path)
make_dir(directory, check_empty=False)
fs = FileSystem.get_file_system(s3_path, 'r')
str_to_file(self.lorem, path)
upload_or_copy(path, s3_path)
stamp = fs.last_modified(s3_path)
self.assertTrue(isinstance(stamp, datetime.datetime))
def save_model_bundle(self):
"""Save a model bundle.
This is a zip file with the model weights in .pth format and a serialized
copy of the LearningConfig, which allows for making predictions in the future.
"""
model_bundle_dir = join(self.tmp_dir, 'model-bundle')
make_dir(model_bundle_dir)
shutil.copyfile(self.last_model_path,
join(model_bundle_dir, 'model.pth'))
shutil.copyfile(self.config_path,
join(model_bundle_dir, 'learner-config.json'))
zipdir(model_bundle_dir, self.model_bundle_path)
def test_file_exists(self):
fs = FileSystem.get_file_system(self.tmp_dir.name, 'r')
path1 = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum.txt')
dir1 = os.path.dirname(path1)
make_dir(dir1, check_empty=False)
str_to_file(self.lorem, path1)
self.assertTrue(fs.file_exists(dir1, include_dir=True))
self.assertTrue(fs.file_exists(path1, include_dir=False))
self.assertFalse(fs.file_exists(dir1, include_dir=False))
self.assertFalse(
fs.file_exists(dir1 + 'NOTPOSSIBLE', include_dir=False))
def test_file_exists(self):
path = os.path.join(self.tmp_dir.name, 'lorem', 'ipsum.txt')
s3_path = 's3://{}/xxx/lorem.txt'.format(self.bucket_name)
s3_path_prefix = 's3://{}/xxx/lorem'.format(self.bucket_name)
s3_directory = 's3://{}/xxx/'.format(self.bucket_name)
make_dir(path, check_empty=False, use_dirname=True)
str_to_file(self.lorem, path)
upload_or_copy(path, s3_path)
self.assertTrue(file_exists(s3_directory, include_dir=True))
self.assertTrue(file_exists(s3_path, include_dir=False))
self.assertFalse(file_exists(s3_path_prefix, include_dir=True))
self.assertFalse(file_exists(s3_directory, include_dir=False))
self.assertFalse(
file_exists(s3_directory + 'NOTPOSSIBLE', include_dir=False))
def save_model_bundle(self):
"""Save a model bundle.
This is a zip file with the model weights in .pth format and a serialized
copy of the LearningConfig, which allows for making predictions in the future.
"""
from rastervision2.pytorch_learner.learner_pipeline_config import (
LearnerPipelineConfig)
model_bundle_dir = join(self.tmp_dir, 'model-bundle')
make_dir(model_bundle_dir)
shutil.copyfile(self.last_model_path,
join(model_bundle_dir, 'model.pth'))
pipeline_cfg = LearnerPipelineConfig(learner=self.cfg)
save_pipeline_config(pipeline_cfg,
join(model_bundle_dir, 'pipeline-config.json'))
zipdir(model_bundle_dir, self.model_bundle_path)
def plot_batch(self, x, y, output_path, z=None):
batch_sz = x.shape[0]
ncols = nrows = math.ceil(math.sqrt(batch_sz))
fig = plt.figure(
constrained_layout=True, figsize=(3 * ncols, 3 * nrows))
grid = gridspec.GridSpec(ncols=ncols, nrows=nrows, figure=fig)
for i in range(batch_sz):
ax = fig.add_subplot(grid[i])
if z is None:
self.plot_xyz(ax, x[i], y[i])
else:
self.plot_xyz(ax, x[i], y[i], z=z[i])
make_dir(output_path, use_dirname=True)
plt.savefig(output_path)
plt.close()
def crop_image(image_uri, window, crop_uri):
im_dataset = rasterio.open(image_uri)
rasterio_window = window.rasterio_format()
im = im_dataset.read(window=rasterio_window)
with tempfile.TemporaryDirectory() as tmp_dir:
crop_path = get_local_path(crop_uri, tmp_dir)
make_dir(crop_path, use_dirname=True)
meta = im_dataset.meta
meta['width'], meta['height'] = window.get_width(), window.get_height()
meta['transform'] = rasterio.windows.transform(rasterio_window,
im_dataset.transform)
with rasterio.open(crop_path, 'w', **meta) as dst:
dst.colorinterp = im_dataset.colorinterp
dst.write(im)
upload_or_copy(crop_path, crop_uri)
def write_sample(self, sample: DataSample):
"""
This writes a training or validation sample to
(train|valid)/{class_name}/{scene_id}-{ind}.png
"""
class_id = sample.labels.get_cell_class_id(sample.window)
# If a chip is not associated with a class, don't
# use it in training data.
if class_id is None:
return
split_name = 'train' if sample.is_train else 'valid'
class_name = self.class_config.names[class_id]
class_dir = join(self.sample_dir, split_name, class_name)
make_dir(class_dir)
chip_path = join(class_dir, '{}-{}.png'.format(sample.scene_id,
self.sample_ind))
save_img(sample.chip, chip_path)
self.sample_ind += 1
def write_sample(self, sample: DataSample):
"""
This writes a training or validation sample to
(train|valid)/img/{scene_id}-{ind}.png and
(train|valid)/labels/{scene_id}-{ind}.png
"""
split_name = 'train' if sample.is_train else 'valid'
label_arr = sample.labels.get_label_arr(sample.window).astype(np.uint8)
img_dir = join(self.sample_dir, split_name, 'img')
labels_dir = join(self.sample_dir, split_name, 'labels')
make_dir(img_dir)
make_dir(labels_dir)
img_path = join(img_dir, '{}-{}.png'.format(sample.scene_id,
self.sample_ind))
labels_path = join(
labels_dir, '{}-{}.png'.format(sample.scene_id, self.sample_ind))
save_img(sample.chip, img_path)
save_img(label_arr, labels_path)
self.sample_ind += 1
def write_sample(self, sample: DataSample):
"""
This writes a training or validation sample to
(train|valid)/img/{scene_id}-{ind}.png and updates
some COCO data structures.
"""
split = 'train' if sample.is_train else 'valid'
split_dir = join(self.sample_dir, split)
img_dir = join(split_dir, 'img')
make_dir(img_dir)
img_fn = '{}-{}.png'.format(sample.scene_id, self.sample_ind)
img_path = join(img_dir, img_fn)
save_img(sample.chip, img_path)
images = self.splits[split]['images']
annotations = self.splits[split]['annotations']
images.append({
'file_name': img_fn,
'id': self.sample_ind,
'height': sample.chip.shape[0],
'width': sample.chip.shape[1]
})
npboxes = sample.labels.get_npboxes()
npboxes = ObjectDetectionLabels.global_to_local(npboxes, sample.window)
for box_ind, (box, class_id) in enumerate(
zip(npboxes, sample.labels.get_class_ids())):
bbox = [box[1], box[0], box[3] - box[1], box[2] - box[0]]
bbox = [int(i) for i in bbox]
annotations.append({
'id': '{}-{}'.format(self.sample_ind, box_ind),
'image_id': self.sample_ind,
'bbox': bbox,
'category_id': int(class_id)
})
self.sample_ind += 1
def __init__(self,
cfg: LearnerConfig,
tmp_dir: str,
model_path: Optional[str] = None):
"""Constructor.
Args:
cfg: configuration
tmp_dir: root of temp dirs
model_path: a local path to model weights. If provided, the model is loaded
and it is assumed that this Learner will be used for prediction only.
"""
self.cfg = cfg
self.tmp_dir = tmp_dir
# TODO make cache dirs configurable
torch_cache_dir = '/opt/data/torch-cache'
os.environ['TORCH_HOME'] = torch_cache_dir
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.data_cache_dir = '/opt/data/data-cache'
make_dir(self.data_cache_dir)
self.model = self.build_model()
self.model.to(self.device)
if model_path is not None:
if isfile(model_path):
self.model.load_state_dict(
torch.load(model_path, map_location=self.device))
else:
raise Exception(
'Model could not be found at {}'.format(model_path))
self.model.eval()
else:
log.info(self.cfg)
# ds = dataset, dl = dataloader
self.train_ds = None
self.train_dl = None
self.valid_ds = None
self.valid_dl = None
self.test_ds = None
self.test_dl = None
if cfg.output_uri.startswith('s3://'):
self.output_dir = get_local_path(cfg.output_uri, tmp_dir)
make_dir(self.output_dir, force_empty=True)
if not cfg.overfit_mode:
self.sync_from_cloud()
else:
self.output_dir = cfg.output_uri
make_dir(self.output_dir)
self.last_model_path = join(self.output_dir, 'last-model.pth')
self.config_path = join(self.output_dir, 'config.json')
self.train_state_path = join(self.output_dir, 'train-state.json')
self.log_path = join(self.output_dir, 'log.csv')
model_bundle_fn = basename(cfg.get_model_bundle_uri())
self.model_bundle_path = join(self.output_dir, model_bundle_fn)
self.metric_names = self.build_metric_names()
json_to_file(self.cfg.dict(), self.config_path)
self.load_init_weights()
self.load_checkpoint()
self.opt = self.build_optimizer()
self.setup_data()
self.start_epoch = self.get_start_epoch()
self.steps_per_epoch = len(
self.train_ds) // self.cfg.solver.batch_sz
self.step_scheduler = self.build_step_scheduler()
self.epoch_scheduler = self.build_epoch_scheduler()
self.setup_tensorboard()
def test_use_dirname(self):
path = os.path.join(self.tmp_dir.name, 'hello', 'hello.txt')
dir = os.path.dirname(path)
make_dir(path, use_dirname=True)
self.assertTrue(os.path.isdir(dir))
def test_default_args(self):
dir = os.path.join(self.tmp_dir.name, 'hello')
make_dir(dir)
self.assertTrue(os.path.isdir(dir))
def test_file_exists_local_false(self):
path = os.path.join(self.tmp_dir.name, 'hello', 'hello.txt')
directory = os.path.dirname(path)
make_dir(directory, check_empty=False)
self.assertFalse(file_exists(path))
