Adding graceful and informative gpu out of memory error handling.

This commit is contained in:
Jason Antic
2020-04-21 18:08:45 -07:00
parent e55c1abc86
commit 696001361a
2 changed files with 15 additions and 5 deletions
+11 -4
View File
@@ -47,10 +47,17 @@ class BaseFilter(IFilter):
x = x.to(self.device)
x.div_(255)
x, y = self.norm((x, x), do_x=True)
result = self.learn.pred_batch(
ds_type=DatasetType.Valid, batch=(x[None], y[None]), reconstruct=True
)
try:
result = self.learn.pred_batch(
ds_type=DatasetType.Valid, batch=(x[None], y[None]), reconstruct=True
)
except RuntimeError as rerr:
if 'memory' not in str(rerr):
raise rerr
print('Warning: render_factor was set too high, and out of memory error resulted. Returning original image.')
return model_image
out = result[0]
out = self.denorm(out.px, do_x=False)
out = image2np(out * 255).astype(np.uint8)
+4 -1
View File
@@ -112,7 +112,10 @@ class ModelImageVisualizer:
else:
self._plot_solo(figsize, render_factor, display_render_factor, result)
return self._save_result_image(path, result)
orig.close()
result_path = self._save_result_image(path, result)
result.close()
return result_path
def _plot_comparison(
self,