data_juicer.ops.common.hawor_func module#
- data_juicer.ops.common.hawor_func.default_config() CfgNode[source]#
Get a yacs CfgNode object with the default config values.
- data_juicer.ops.common.hawor_func.get_config(config_file: str, merge: bool = True, update_cachedir: bool = False) CfgNode[source]#
Read a config file and optionally merge it with the default config file. :param config_file: Path to config file. :type config_file: str :param merge: Whether to merge with the default config or not. :type merge: bool
- Returns:
Config as a yacs CfgNode object.
- Return type:
CfgNode
- data_juicer.ops.common.hawor_func.normalization_layer(norm: str | None, dim: int, norm_cond_dim: int = -1)[source]#
- class data_juicer.ops.common.hawor_func.AdaptiveLayerNorm1D(data_dim: int, norm_cond_dim: int)[source]#
Bases:
Module- __init__(data_dim: int, norm_cond_dim: int)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x: Tensor, t: Tensor) Tensor[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.FrequencyEmbedder(num_frequencies, max_freq_log2)[source]#
Bases:
Module- __init__(num_frequencies, max_freq_log2)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- data_juicer.ops.common.hawor_func.get_transform(center, scale, res, rot=0)[source]#
Generate transformation matrix.
- data_juicer.ops.common.hawor_func.transform(pt, center, scale, res, invert=0, rot=0, asint=True)[source]#
Transform pixel location to different reference.
- data_juicer.ops.common.hawor_func.crop(img, center, scale, res, rot=0)[source]#
Crop image according to the supplied bounding box.
- data_juicer.ops.common.hawor_func.perspective_projection(points, rotation, translation, focal_length, camera_center, distortion=None)[source]#
This function computes the perspective projection of a set of points. Input:
points (bs, N, 3): 3D points rotation (bs, 3, 3): Camera rotation translation (bs, 3): Camera translation focal_length (bs,) or scalar: Focal length camera_center (bs, 2): Camera center
- data_juicer.ops.common.hawor_func.rot9d_to_rotmat(x)[source]#
Convert 9D rotation representation to 3x3 rotation matrix. Based on Levinson et al., “An Analysis of SVD for Deep Rotation Estimation” Input:
(B,9) or (B,J*9) Batch of 9D rotation (interpreted as 3x3 est rotmat)
- Output:
(B,3,3) or (B*J,3,3) Batch of corresponding rotation matrices
- data_juicer.ops.common.hawor_func.quat_to_rotmat(quat)[source]#
Convert quaternion coefficients to rotation matrix. :param quat: size = [B, 4] 4 <===>(w, x, y, z)
- Returns:
Rotation matrix corresponding to the quaternion – size = [B, 3, 3]
- data_juicer.ops.common.hawor_func.rot6d_to_rotmat(x: Tensor) Tensor[source]#
Convert 6D rotation representation to 3x3 rotation matrix. Based on Zhou et al., “On the Continuity of Rotation Representations in Neural Networks”, CVPR 2019 :param x: (B,6) Batch of 6-D rotation representations. :type x: torch.Tensor
- Returns:
Batch of corresponding rotation matrices with shape (B,3,3).
- Return type:
torch.Tensor
- data_juicer.ops.common.hawor_func.rotmat_to_rot6d(rotmat)[source]#
Inverse function of the above. Input:
(B,3,3) Batch of corresponding rotation matrices
- Output:
(B,6) Batch of 6-D rotation representations
- data_juicer.ops.common.hawor_func.rotation_matrix_to_angle_axis(rotation_matrix)[source]#
This function is borrowed from kornia/kornia
Convert 3x4 rotation matrix to Rodrigues vector
- Parameters:
rotation_matrix (Tensor) – rotation matrix.
- Returns:
Rodrigues vector transformation.
- Return type:
Tensor
- Shape:
Input: \((N, 3, 4)\)
Output: \((N, 3)\)
Example
>>> input = torch.rand(2, 3, 4) # Nx4x4 >>> output = tgm.rotation_matrix_to_angle_axis(input) # Nx3
- data_juicer.ops.common.hawor_func.quaternion_to_angle_axis(quaternion: Tensor) Tensor[source]#
This function is borrowed from kornia/kornia
Convert quaternion vector to angle axis of rotation.
Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
- Parameters:
quaternion (torch.Tensor) – tensor with quaternions.
- Returns:
tensor with angle axis of rotation.
- Return type:
torch.Tensor
- Shape:
Input: \((*, 4)\) where * means, any number of dimensions
Output: \((*, 3)\)
Example
>>> quaternion = torch.rand(2, 4) # Nx4 >>> angle_axis = tgm.quaternion_to_angle_axis(quaternion) # Nx3
- data_juicer.ops.common.hawor_func.rotation_matrix_to_quaternion(rotation_matrix, eps=1e-06)[source]#
This function is borrowed from kornia/kornia Convert rotation matrix to 4d quaternion vector This algorithm is based on algorithm described in KieranWynn/pyquaternion
:param rotation_matrix (N, 3, 3)
- data_juicer.ops.common.hawor_func.estimate_translation_np(S, joints_2d, joints_conf, focal_length=5000.0, img_size=224.0)[source]#
This function is borrowed from nkolot/SPIN
Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d. Input:
S: (25, 3) 3D joint locations joints: (25, 3) 2D joint locations and confidence
- Returns:
(3,) camera translation vector
- data_juicer.ops.common.hawor_func.estimate_translation(S, joints_2d, focal_length=5000.0, img_size=224.0)[source]#
Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d. Input:
S: (B, 49, 3) 3D joint locations joints: (B, 49, 3) 2D joint locations and confidence
- Returns:
(B, 3) camera translation vectors
- data_juicer.ops.common.hawor_func.get_keypoints_rectangle(keypoints: array, threshold: float) Tuple[float, float, float][source]#
Compute rectangle enclosing keypoints above the threshold. :param keypoints: Keypoint array of shape (N, 3). :type keypoints: np.array :param threshold: Confidence visualization threshold. :type threshold: float
- Returns:
Rectangle width, height and area.
- Return type:
Tuple[float, float, float]
- data_juicer.ops.common.hawor_func.render_keypoints(img: array, keypoints: array, pairs: List, colors: List, thickness_circle_ratio: float, thickness_line_ratio_wrt_circle: float, pose_scales: List, threshold: float = 0.1, alpha: float = 1.0) array[source]#
Render keypoints on input image. :param img: Input image of shape (H, W, 3) with pixel values in the [0,255] range. :type img: np.array :param keypoints: Keypoint array of shape (N, 3). :type keypoints: np.array :param pairs: List of keypoint pairs per limb. :type pairs: List :param colors: (List): List of colors per keypoint. :param thickness_circle_ratio: Circle thickness ratio. :type thickness_circle_ratio: float :param thickness_line_ratio_wrt_circle: Line thickness ratio wrt the circle. :type thickness_line_ratio_wrt_circle: float :param pose_scales: List of pose scales. :type pose_scales: List :param threshold: Only visualize keypoints with confidence above the threshold. :type threshold: float
- Returns:
Image of shape (H, W, 3) with keypoints drawn on top of the original image.
- Return type:
(np.array)
- data_juicer.ops.common.hawor_func.render_hand_keypoints(img, right_hand_keypoints, threshold=0.1, use_confidence=False, map_fn=<function <lambda>>, alpha=1.0)[source]#
- data_juicer.ops.common.hawor_func.render_hand_landmarks(img, right_hand_keypoints, threshold=0.1, use_confidence=False, map_fn=<function <lambda>>, alpha=1.0)[source]#
- data_juicer.ops.common.hawor_func.render_body_keypoints(img: array, body_keypoints: array) array[source]#
Render OpenPose body keypoints on input image. :param img: Input image of shape (H, W, 3) with pixel values in the [0,255] range. :type img: np.array :param body_keypoints: Keypoint array of shape (N, 3); 3 <====> (x, y, confidence). :type body_keypoints: np.array
- Returns:
Image of shape (H, W, 3) with keypoints drawn on top of the original image.
- Return type:
(np.array)
- data_juicer.ops.common.hawor_func.render_openpose(img: array, hand_keypoints: array) array[source]#
Render keypoints in the OpenPose format on input image. :param img: Input image of shape (H, W, 3) with pixel values in the [0,255] range. :type img: np.array :param body_keypoints: Keypoint array of shape (N, 3); 3 <====> (x, y, confidence). :type body_keypoints: np.array
- Returns:
Image of shape (H, W, 3) with keypoints drawn on top of the original image.
- Return type:
(np.array)
- data_juicer.ops.common.hawor_func.render_openpose_landmarks(img: array, hand_keypoints: array) array[source]#
Render keypoints in the OpenPose format on input image. :param img: Input image of shape (H, W, 3) with pixel values in the [0,255] range. :type img: np.array :param body_keypoints: Keypoint array of shape (N, 3); 3 <====> (x, y, confidence). :type body_keypoints: np.array
- Returns:
Image of shape (H, W, 3) with keypoints drawn on top of the original image.
- Return type:
(np.array)
- data_juicer.ops.common.hawor_func.get_pylogger(name='data_juicer.ops.common.hawor_func') Logger[source]#
Initializes multi-GPU-friendly python command line logger.
- class data_juicer.ops.common.hawor_func.PreNorm(dim: int, fn: Callable, norm: str = 'layer', norm_cond_dim: int = -1)[source]#
Bases:
Module- __init__(dim: int, fn: Callable, norm: str = 'layer', norm_cond_dim: int = -1)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x: Tensor, *args, **kwargs)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.FeedForward(dim, hidden_dim, dropout=0.0)[source]#
Bases:
Module- __init__(dim, hidden_dim, dropout=0.0)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.Attention(dim, heads=8, dim_head=64, dropout=0.0)[source]#
Bases:
Module- __init__(dim, heads=8, dim_head=64, dropout=0.0)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.CrossAttention(dim, context_dim=None, heads=8, dim_head=64, dropout=0.0)[source]#
Bases:
Module- __init__(dim, context_dim=None, heads=8, dim_head=64, dropout=0.0)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x, context=None)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.Transformer(dim: int, depth: int, heads: int, dim_head: int, mlp_dim: int, dropout: float = 0.0, norm: str = 'layer', norm_cond_dim: int = -1)[source]#
Bases:
Module- __init__(dim: int, depth: int, heads: int, dim_head: int, mlp_dim: int, dropout: float = 0.0, norm: str = 'layer', norm_cond_dim: int = -1)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x: Tensor, *args)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.TransformerCrossAttn(dim: int, depth: int, heads: int, dim_head: int, mlp_dim: int, dropout: float = 0.0, norm: str = 'layer', norm_cond_dim: int = -1, context_dim: int | None = None)[source]#
Bases:
Module- __init__(dim: int, depth: int, heads: int, dim_head: int, mlp_dim: int, dropout: float = 0.0, norm: str = 'layer', norm_cond_dim: int = -1, context_dim: int | None = None)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x: Tensor, *args, context=None, context_list=None)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.DropTokenDropout(p: float = 0.1)[source]#
Bases:
Module- __init__(p: float = 0.1)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x: Tensor)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.ZeroTokenDropout(p: float = 0.1)[source]#
Bases:
Module- __init__(p: float = 0.1)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x: Tensor)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.TransformerEncoder(num_tokens: int, token_dim: int, dim: int, depth: int, heads: int, mlp_dim: int, dim_head: int = 64, dropout: float = 0.0, emb_dropout: float = 0.0, emb_dropout_type: str = 'drop', emb_dropout_loc: str = 'token', norm: str = 'layer', norm_cond_dim: int = -1, token_pe_numfreq: int = -1)[source]#
Bases:
Module- __init__(num_tokens: int, token_dim: int, dim: int, depth: int, heads: int, mlp_dim: int, dim_head: int = 64, dropout: float = 0.0, emb_dropout: float = 0.0, emb_dropout_type: str = 'drop', emb_dropout_loc: str = 'token', norm: str = 'layer', norm_cond_dim: int = -1, token_pe_numfreq: int = -1)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(inp: Tensor, *args, **kwargs)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.TransformerDecoder(num_tokens: int, token_dim: int, dim: int, depth: int, heads: int, mlp_dim: int, dim_head: int = 64, dropout: float = 0.0, emb_dropout: float = 0.0, emb_dropout_type: str = 'drop', norm: str = 'layer', norm_cond_dim: int = -1, context_dim: int | None = None, skip_token_embedding: bool = False)[source]#
Bases:
Module- __init__(num_tokens: int, token_dim: int, dim: int, depth: int, heads: int, mlp_dim: int, dim_head: int = 64, dropout: float = 0.0, emb_dropout: float = 0.0, emb_dropout_type: str = 'drop', norm: str = 'layer', norm_cond_dim: int = -1, context_dim: int | None = None, skip_token_embedding: bool = False)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(inp: Tensor, *args, context=None, context_list=None)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.MANOTransformerDecoderHead(cfg)[source]#
Bases:
ModuleHMR2 Cross-attention based SMPL Transformer decoder
- forward(x, **kwargs)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.temporal_attention(in_dim=1280, out_dim=1280, hdim=512, nlayer=6, nhead=4, residual=False)[source]#
Bases:
Module- __init__(in_dim=1280, out_dim=1280, hdim=512, nlayer=6, nhead=4, residual=False)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.PositionalEncoding(d_model, dropout=0.1, max_len=100)[source]#
Bases:
Module- __init__(d_model, dropout=0.1, max_len=100)[source]#
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- forward(x)[source]#
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
- class data_juicer.ops.common.hawor_func.TrackDatasetEval(imgfiles, boxes, crop_size=256, dilate=1.0, img_focal=None, img_center=None, normalization=True, item_idx=0, do_flip=False)[source]#
Bases:
DatasetTrack Dataset Class - Load images/crops of the tracked boxes.
- data_juicer.ops.common.hawor_func.dataset_config() CfgNode[source]#
Get dataset config file :returns: Dataset config as a yacs CfgNode object. :rtype: CfgNode
- data_juicer.ops.common.hawor_func.parse_chunks(frame, boxes, min_len=16)[source]#
If a track disappear in the middle, we separate it to different segments to estimate the HPS independently. If a segment is less than 16 frames, we get rid of it for now.
- data_juicer.ops.common.hawor_func.aa_to_rotmat(theta: Tensor)[source]#
Convert axis-angle representation to rotation matrix. Works by first converting it to a quaternion. :param theta: Tensor of shape (B, 3) containing axis-angle representations. :type theta: torch.Tensor
- Returns:
Corresponding rotation matrices with shape (B, 3, 3).
- Return type:
torch.Tensor
- class data_juicer.ops.common.hawor_func.MANO(*args, joint_regressor_extra: str | None = None, **kwargs)[source]#
Bases:
MANOLayer- __init__(*args, joint_regressor_extra: str | None = None, **kwargs)[source]#
Extension of the official MANO implementation to support more joints. :param Same as MANOLayer.: :param joint_regressor_extra: Path to extra joint regressor. :type joint_regressor_extra: str
- data_juicer.ops.common.hawor_func.run_mano(trans, root_orient, hand_pose, is_right=None, betas=None, use_cuda=True)[source]#
Forward pass of the SMPL model and populates pred_data accordingly with joints3d, verts3d, points3d.
trans : B x T x 3 root_orient : B x T x 3 body_pose : B x T x J*3 betas : (optional) B x D
- data_juicer.ops.common.hawor_func.run_mano_left(trans, root_orient, hand_pose, is_right=None, betas=None, use_cuda=True, fix_shapedirs=True)[source]#
Forward pass of the SMPL model and populates pred_data accordingly with joints3d, verts3d, points3d.
trans : B x T x 3 root_orient : B x T x 3 body_pose : B x T x J*3 betas : (optional) B x D
- data_juicer.ops.common.hawor_func.run_mano_twohands(init_trans, init_rot, init_hand_pose, is_right, init_betas, use_cuda=True, fix_shapedirs=True)[source]#
- data_juicer.ops.common.hawor_func.quaternion_mul(q0, q1)[source]#
EXPECTS WXYZ :param q0 (, 4) :param q1 (, 4)
- data_juicer.ops.common.hawor_func.quaternion_inverse(q, eps=1e-08)[source]#
EXPECTS WXYZ :param q (*, 4)
- data_juicer.ops.common.hawor_func.quaternion_slerp(t, q0, q1, eps=1e-08)[source]#
:param t (, 1) must be between 0 and 1 :param q0 (, 4) :param q1 (*, 4)
- data_juicer.ops.common.hawor_func.angle_axis_to_rotation_matrix(angle_axis)[source]#
:param angle_axis (, 3) return (, 3, 3)
- data_juicer.ops.common.hawor_func.quaternion_to_rotation_matrix(quaternion)[source]#
Convert a quaternion to a rotation matrix. Taken from kornia/kornia, based on matthew-brett/transforms3d tensorflow/graphics :param quaternion (N, 4) expects WXYZ order returns rotation matrix (N, 3, 3)
- data_juicer.ops.common.hawor_func.angle_axis_to_quaternion(angle_axis)[source]#
This function is borrowed from kornia/kornia Convert angle axis to quaternion in WXYZ order :param angle_axis (, 3) :returns quaternion (, 4) WXYZ order
- class data_juicer.ops.common.hawor_func.HAWOR(cfg: CfgNode)[source]#
Bases:
LightningModule- __init__(cfg: CfgNode)[source]#
Setup HAWOR model :param cfg: Config file as a yacs CfgNode :type cfg: CfgNode
- configure_optimizers() Optimizer[source]#
Setup model and distriminator Optimizers :returns: Model and discriminator optimizers :rtype: Tuple[torch.optim.Optimizer, torch.optim.Optimizer]
- forward_step(batch: Dict, train: bool = False) Dict[source]#
Run a forward step of the network :param batch: Dictionary containing batch data :type batch: Dict :param train: Flag indicating whether it is training or validation mode :type train: bool
- Returns:
Dictionary containing the regression output
- Return type:
Dict
- compute_loss(batch: Dict, output: Dict, train: bool = True) Tensor[source]#
Compute losses given the input batch and the regression output :param batch: Dictionary containing batch data :type batch: Dict :param output: Dictionary containing the regression output :type output: Dict :param train: Flag indicating whether it is training or validation mode :type train: bool
- Returns:
Total loss for current batch
- Return type:
torch.Tensor
- tensorboard_logging(batch: Dict, output: Dict, step_count: int, train: bool = True, write_to_summary_writer: bool = True, render_log: bool = True) None[source]#
Log results to Tensorboard :param batch: Dictionary containing batch data :type batch: Dict :param output: Dictionary containing the regression output :type output: Dict :param step_count: Global training step count :type step_count: int :param train: Flag indicating whether it is training or validation mode :type train: bool
- forward(batch: Dict) Dict[source]#
Run a forward step of the network in val mode :param batch: Dictionary containing batch data :type batch: Dict
- Returns:
Dictionary containing the regression output
- Return type:
Dict
- training_step(joint_batch: Dict, batch_idx: int) Dict[source]#
Run a full training step :param joint_batch: Dictionary containing image and mocap batch data :type joint_batch: Dict :param batch_idx: Unused. :type batch_idx: int :param batch_idx: Unused. :type batch_idx: torch.Tensor
- Returns:
Dictionary containing regression output.
- Return type:
Dict
- validation_step(batch: Dict, batch_idx: int, dataloader_idx=0) Dict[source]#
Run a validation step and log to Tensorboard :param batch: Dictionary containing batch data :type batch: Dict :param batch_idx: Unused. :type batch_idx: int
- Returns:
Dictionary containing regression output.
- Return type:
Dict