diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index c2de9ea89..4da8c07c6 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -20,6 +20,7 @@ from __future__ import annotations import functools import logging import os +import time from pathlib import Path from typing import Any, Dict, List @@ -120,7 +121,6 @@ class Trainer: with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME): num_iterations = self.config.trainer.max_num_iterations for step in range(self.start_step, self.start_step + num_iterations): - # if the visualizer used, the rendering of the visualizer will be included in the iteration train time with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as train_t: @@ -138,7 +138,14 @@ class Trainer: with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as vis_t: num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch - self.visualizer_state.update_scene(step, self.pipeline.model, num_rays_per_batch) + try: + self.visualizer_state.update_scene(step, self.pipeline.model, num_rays_per_batch) + except RuntimeError: + time.sleep(0.03) # sleep to allow buffer to reset + assert self.visualizer_state.vis is not None + self.visualizer_state.vis["renderingState/log_errors"].write( + "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing." + ) if step != 0 and step % self.config.logging.steps_per_log == 0: writer.put_dict(name="Loss/train-loss_metrics_dict", scalar_dict=loss_metric_dict, step=step) diff --git a/nerfactory/viewer/app/src/modules/LogPanel/LogPanel.jsx b/nerfactory/viewer/app/src/modules/LogPanel/LogPanel.jsx new file mode 100644 index 000000000..b85b4ebe5 --- /dev/null +++ b/nerfactory/viewer/app/src/modules/LogPanel/LogPanel.jsx @@ -0,0 +1,91 @@ +import { useContext, useEffect } from 'react'; +import { useDispatch, useSelector } from 'react-redux'; + +import { WebSocketContext } from '../WebSocket/WebSocket'; + +const msgpack = require('msgpack-lite'); + +export function LogPanel() { + const websocket = useContext(WebSocketContext).socket; + const dispatch = useDispatch(); + const gpu_oom_error_msg = "GPU out of memory"; + const resolved_msg = "resolved"; + let local_error = resolved_msg; + // connection status indicators + + const set_max_train_util = () => { + if (websocket.readyState === WebSocket.OPEN) { + dispatch({ + type: 'write', + path: 'renderingState/targetTrainUtil', + data: 0.9, + }); + const cmd = 'write'; + const path = 'renderingState/targetTrainUtil'; + const data = { + type: cmd, + path, + data: 0.9, + }; + const message = msgpack.encode(data); + websocket.send(message); + } + }; + + const set_small_resolution = () => { + if (websocket.readyState === WebSocket.OPEN) { + dispatch({ + type: 'write', + path: 'renderingState/maxResolution', + data: 512, + }); + const cmd = 'write'; + const path = 'renderingState/maxResolution'; + const data = { + type: cmd, + path, + data: 512, + }; + const message = msgpack.encode(data); + websocket.send(message); + } + }; + + const set_log_message = () => { + if (websocket.readyState === WebSocket.OPEN) { + dispatch({ + type: 'write', + path: 'renderingState/log_errors', + data: resolved_msg, + }); + const cmd = 'write'; + const path = 'renderingState/log_errors'; + const data = { + type: cmd, + path, + data: resolved_msg, + }; + const message = msgpack.encode(data); + websocket.send(message); + } + }; + + const check_error = useSelector( + (state) => { + local_error = state.renderingState.log_errors; + if(local_error.includes(gpu_oom_error_msg)){ + console.log(local_error); + set_log_message(); + set_small_resolution(); + set_max_train_util(); + } + } + ); + + useEffect(() => {}, [ + check_error, + local_error, + ]); + + return null; +} diff --git a/nerfactory/viewer/app/src/modules/SidePanel/SidePanel.jsx b/nerfactory/viewer/app/src/modules/SidePanel/SidePanel.jsx index 041583e9d..7eda3e863 100644 --- a/nerfactory/viewer/app/src/modules/SidePanel/SidePanel.jsx +++ b/nerfactory/viewer/app/src/modules/SidePanel/SidePanel.jsx @@ -20,6 +20,7 @@ import SceneNode from '../../SceneNode'; import LevaTheme from '../ConfigPanel/leva_theme.json'; import CameraPanel from './CameraPanel'; import { RenderControls } from '../ConfigPanel/ConfigPanel'; +import { LogPanel } from '../LogPanel/LogPanel'; interface TabPanelProps { children: React.ReactNode; @@ -182,7 +183,7 @@ export function BasicTabs(props: BasicTabsProps) { - TODO: add a logging panel + TODO: Something maybe? diff --git a/nerfactory/viewer/app/src/reducer.js b/nerfactory/viewer/app/src/reducer.js index 8e3ae8ce8..c4d1698a4 100644 --- a/nerfactory/viewer/app/src/reducer.js +++ b/nerfactory/viewer/app/src/reducer.js @@ -33,6 +33,7 @@ const initialState = { eval_fps: 'Paused', train_eta: 'Paused', vis_train_ratio: 'Paused', + log_errors: '', }, // the scene state sceneState: { diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 56ed25046..5fa261245 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -20,6 +20,7 @@ import logging import os import sys import threading +import time from pathlib import Path from typing import Any, Dict, Optional @@ -106,8 +107,9 @@ class RenderThread(threading.Thread): outputs = None try: with SetTrace(self.state.check_interrupt): - outputs = self.graph.get_outputs_for_camera_ray_bundle(self.camera_ray_bundle) - except IOChangeException as e: + with torch.no_grad(): + outputs = self.graph.get_outputs_for_camera_ray_bundle(self.camera_ray_bundle) + except Exception as e: # pylint: disable=broad-except self.exc = e if outputs: @@ -119,7 +121,6 @@ class RenderThread(threading.Thread): def join(self, timeout=None): threading.Thread.join(self) - torch.cuda.empty_cache() if self.exc: raise self.exc @@ -544,7 +545,14 @@ class VisualizerState: self.prev_colormap_type = colormap_type # Calculate camera pose and intrinsics - image_height = self._calculate_image_height(camera_object, is_training) + try: + image_height = self._calculate_image_height(camera_object, is_training) + except ZeroDivisionError as e: + self.vis["renderingState/log_errors"].write("Error: Screen too small; no rays intersecting scene.") + time.sleep(0.03) # sleep to allow buffer to reset + print(f"Error: {e}") + return + if image_height is None: return @@ -585,8 +593,17 @@ class VisualizerState: try: render_thread.join() check_thread.join() - except Exception: # pylint: disable=broad-except + except IOChangeException: pass + except RuntimeError as e: + del camera_ray_bundle + torch.cuda.empty_cache() + time.sleep(0.5) # sleep to allow buffer to reset + self.vis["renderingState/log_errors"].write( + "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing." + ) + print(f"Error: {e}") + graph.train() outputs = render_thread.vis_outputs if outputs is not None: