def save_model(self,
model,
name,
n_iter=None,
iter_format="{:05d}",
prefix=False):
"""
Saves a pytorch model in the model directory of the experiment folder
Args:
model: The model to be stored
name: The file name of the model file
n_iter: The iteration number, formatted with the iter_format and added to the model name (if not None)
iter_format: The format string, which indicates how n_iter will be formated as a string
prefix: If True, the formated n_iter will be appended as a prefix, otherwise as a suffix
"""
if n_iter is not None:
name = name_and_iter_to_filename(name,
n_iter,
".pth",
iter_format=iter_format,
prefix=prefix)
if not name.endswith(".pth"):
name += ".pth"
self.save_model_static(model=model,
model_dir=self.checkpoint_dir,
name=name)
def save_checkpoint(self,
name,
n_iter=None,
iter_format="{:05d}",
prefix=False,
**kwargs):
"""
Saves a checkpoint in the checkpoint directory of the experiment folder
Args:
name: The file name of the checkpoint file
n_iter: The iteration number, formatted with the iter_format and added to the checkpoint name (if not None)
iter_format: The format string, which indicates how n_iter will be formated as a string
prefix: If True, the formated n_iter will be appended as a prefix, otherwise as a suffix
**kwargs: dict which is actually saved (key=name, value=variable to be stored)
"""
if n_iter is not None:
name = name_and_iter_to_filename(name,
n_iter,
".pth.tar",
iter_format=iter_format,
prefix=prefix)
if not name.endswith(".pth.tar"):
name += ".pth.tar"
self.save_checkpoint_static(self.checkpoint_dir, name=name, **kwargs)
def show_image_grid_heatmap(self,
heatmap,
background=None,
ratio=0.3,
normalize=True,
colormap=cm.jet,
name="heatmap",
n_iter=None,
prefix=False,
iter_format="{:05d}",
image_args=None,
**kwargs):
"""
Creates heat map from the given map and if given combines it with the background and then
displays results with as image grid.
Args:
heatmap: 4d- tensor (N, C, H, W) to be converted to a heatmap
background: 4d- tensor (N, C, H, W) background/ context of the heatmap (to be underlayed)
name: The name of the window
ratio: The ratio to mix the map with the background (0 = only background, 1 = only map)
n_iter: The iteration number, formatted with the iter_format and added to the model name (if not None)
iter_format: The format string, which indicates how n_iter will be formated as a string
prefix: If True, the formated n_iter will be appended as a prefix, otherwise as a suffix
image_args: Arguments for the tensorvision save image method
"""
if image_args is None:
image_args = {}
if "normalize" not in image_args:
image_args["normalize"] = normalize
if n_iter is not None:
name = name_and_iter_to_filename(name=name,
n_iter=n_iter,
ending=".png",
iter_format=iter_format,
prefix=prefix)
elif not name.endswith(".png"):
name += ".png"
file_name = os.path.join(self.img_dir, name)
map_grid = np_make_grid(
heatmap, normalize=normalize) # map_grid.shape is (3, X, Y)
if heatmap.shape[1] != 3:
map_ = colormap(map_grid[0])[..., :-1].transpose(2, 0, 1)
else: # heatmap was already RGB, so don't apply colormap
map_ = map_grid
if background is not None:
img_grid = np_make_grid(background, **image_args)
fuse_img = (1.0 - ratio) * img_grid + ratio * map_
else:
fuse_img = map_
fuse_img = np.clip(fuse_img * 255, a_min=0, a_max=255).astype(np.uint8)
imwrite(file_name, fuse_img.transpose(1, 2, 0))
def save_image_static(image_dir, tensor, name, n_iter=None, iter_format="{:05d}", prefix=False, image_args=None):
"""
Saves an image tensor in an image directory
Args:
image_dir: Directory to save the image in
tensor: Tensor containing the image
name: file-name of the image file
n_iter: The iteration number, formatted with the iter_format and added to the model name (if not None)
iter_format: The format string, which indicates how n_iter will be formated as a string
prefix: If True, the formated n_iter will be appended as a prefix, otherwise as a suffix
image_args: Arguments for the tensorvision save image method
"""
if isinstance(tensor, np.ndarray):
tensor = torch.from_numpy(tensor)
if image_args is None:
image_args = {}
if n_iter is not None:
name = name_and_iter_to_filename(name=name, n_iter=n_iter, ending=".png", iter_format=iter_format,
prefix=prefix)
elif not name.endswith(".png"):
name = name + ".png"
img_file = os.path.join(image_dir, name)
os.makedirs(os.path.dirname(img_file), exist_ok=True)
tv_save_image(tensor=tensor, filename=img_file, **image_args)
def save_image_grid_static(image_dir,
tensor,
name,
n_iter=None,
prefix=False,
iter_format="{:05d}",
image_args=None):
"""
Saves images of a 4d- tensor (N, C, H, W) as a image grid into an image file in a given directory
Args:
image_dir: Directory to save the image in
tensor: 4d- tensor (N, C, H, W)
name: file-name of the image file
n_iter: The iteration number, formatted with the iter_format and added to the model name (if not None)
iter_format: The format string, which indicates how n_iter will be formated as a string
prefix: If True, the formated n_iter will be appended as a prefix, otherwise as a suffix
image_args: Arguments for the tensorvision save image method
"""
if isinstance(tensor, np.ndarray):
tensor = torch.tensor(tensor)
if not (tensor.size(1) == 1 or tensor.size(1) == 3):
warnings.warn(
"The 1. dimension (channel) has to be either 1 (gray) or 3 (rgb), taking the first "
"dimension now !!!")
tensor = tensor[:, 0:1, ]
if n_iter is not None:
name = name_and_iter_to_filename(name=name,
n_iter=n_iter,
ending=".png",
iter_format=iter_format,
prefix=prefix)
elif not name.endswith(".png"):
name += ".png"
img_file = os.path.join(image_dir, name)
if image_args is None:
image_args = {}
os.makedirs(os.path.dirname(img_file), exist_ok=True)
tv_save_image(tensor=tensor, filename=img_file, **image_args)
def show_image_grid_heatmap(self, heatmap, background=None, ratio=0.3, normalize=True,
colormap=cv2.COLORMAP_JET, name="heatmap", n_iter=None,
prefix=False, iter_format="{:05d}", image_args=None, **kwargs):
"""
Creates heat map from the given map and if given combines it with the background and then
displays results with as image grid.
Args:
heatmap: 4d- tensor (N, C, H, W) to be converted to a heatmap
background: 4d- tensor (N, C, H, W) background/ context of the heatmap (to be underlayed)
name: The name of the window
ratio: The ratio to mix the map with the background (0 = only background, 1 = only map)
n_iter: The iteration number, formatted with the iter_format and added to the model name (if not None)
iter_format: The format string, which indicates how n_iter will be formated as a string
prefix: If True, the formated n_iter will be appended as a prefix, otherwise as a suffix
image_args: Arguments for the tensorvision save image method
"""
if image_args is None: image_args = {}
if n_iter is not None:
name = name_and_iter_to_filename(name=name, n_iter=n_iter, ending=".png", iter_format=iter_format,
prefix=prefix)
elif not name.endswith(".png"):
name += ".png"
file_name = os.path.join(self.img_dir, name)
map_grid = np_make_grid(heatmap, normalize=normalize)
map_ = np.clip(map_grid * 255, a_min=0, a_max=255)
map_ = map_.astype(np.uint8)
map_ = cv2.applyColorMap(map_.transpose(1, 2, 0), colormap=colormap)
map_ = cv2.cvtColor(map_, cv2.COLOR_BGR2RGB)
map_ = map_.transpose(2, 0, 1)
fuse_img = map_
if background is not None:
img_grid = np_make_grid(background, **image_args)
image = np.clip(img_grid * 255, a_min=0, a_max=255)
image = image.astype(np.uint8)
fuse_img = (1.0 - ratio) * image + ratio * map_
imsave(file_name, fuse_img.transpose(1, 2, 0))
def load_checkpoint(self,
name="checkpoint",
save_types=("model", "optimizer", "simple", "th_vars",
"results"),
n_iter=None,
iter_format="{:05d}",
prefix=False,
path=None):
"""
Loads a checkpoint and restores the experiment.
Make sure you have your torch stuff already on the right devices beforehand,
otherwise this could lead to errors e.g. when making a optimizer step
(and for some reason the Adam states are not already on the GPU:
https://discuss.pytorch.org/t/loading-a-saved-model-for-continue-training/17244/3 )
Args:
name (str): The name of the checkpoint file
save_types (list or tuple): What kind of member variables should be loaded? Choices are:
"model" <-- Pytorch models,
"optimizer" <-- Optimizers,
"simple" <-- Simple python variables (basic types and lists/tuples),
"th_vars" <-- torch tensors,
"results" <-- The result dict
n_iter (int): Number of iterations. Together with the name, defined by the iter_format,
a file name will be created and searched for.
iter_format (str): Defines how the name and the n_iter will be combined.
prefix (bool): If True, the formatted n_iter will be prepended, otherwise appended.
path (str): If no path is given then it will take the current experiment dir and formatted
name, otherwise it will simply use the path and the formatted name to define the
checkpoint file.
"""
if self.elog is None:
return
model_dict = {}
optimizer_dict = {}
simple_dict = {}
th_vars_dict = {}
results_dict = {}
if "model" in save_types:
model_dict = self.get_pytorch_modules()
if "optimizer" in save_types:
optimizer_dict = self.get_pytorch_optimizers()
if "simple" in save_types:
simple_dict = self.get_simple_variables()
if "th_vars" in save_types:
th_vars_dict = self.get_pytorch_variables()
if "results" in save_types:
results_dict = {"results": self.results}
checkpoint_dict = {
**model_dict,
**optimizer_dict,
**simple_dict,
**th_vars_dict,
**results_dict
}
if n_iter is not None:
name = name_and_iter_to_filename(name,
n_iter,
".pth.tar",
iter_format=iter_format,
prefix=prefix)
if path is None:
restore_dict = self.elog.load_checkpoint(name=name,
**checkpoint_dict)
else:
checkpoint_path = os.path.join(path, name)
if checkpoint_path.endswith("/"):
checkpoint_path = checkpoint_path[:-1]
restore_dict = self.elog.load_checkpoint_static(
checkpoint_file=checkpoint_path, **checkpoint_dict)
self.update_attributes(restore_dict)
