RuntimeError: einsum() dimension mismatch in apply_deformation_transfer function

[ajfabian]

Issue Description

I encountered a RuntimeError while running the transfer_model script. The error occurs in the apply_deformation_transfer function within the def_transfer.py file. Specifically, the torch.einsum operation fails due to a dimension mismatch.

Error Traceback

  0%|                                                                                                                                                          | 0/1 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "/home/ubuntu/miniconda3/envs/ml-monorepo/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/ubuntu/miniconda3/envs/ml-monorepo/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/ubuntu/app/ml-projects-monorepo/smplx/transfer_model/__main__.py", line 104, in <module>
    main()
  File "/home/ubuntu/app/ml-projects-monorepo/smplx/transfer_model/__main__.py", line 84, in main
    var_dict = run_fitting(
  File "/home/ubuntu/app/ml-projects-monorepo/smplx/transfer_model/transfer_model.py", line 280, in run_fitting
    def_vertices = apply_deformation_transfer(def_matrix, vertices, faces)
  File "/home/ubuntu/app/ml-projects-monorepo/smplx/transfer_model/utils/def_transfer.py", line 78, in apply_deformation_transfer
    def_vertices = torch.einsum('mn,bni->bmi', [def_matrix, vertices])
  File "/home/ubuntu/miniconda3/envs/ml-monorepo/lib/python3.8/site-packages/torch/functional.py", line 381, in einsum
    return einsum(equation, *_operands)
  File "/home/ubuntu/miniconda3/envs/ml-monorepo/lib/python3.8/site-packages/torch/functional.py", line 386, in einsum
    return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
RuntimeError: einsum(): subscript n has size 11437 for operand 1 which does not broadcast with previously seen size 10475

Steps to Reproduce

1. Run the transfer_model script with the current configuration.
   python -m transfer_model --exp-cfg config_files/smplx2smpl.yaml
2. Observe the RuntimeError in the console output.

Expected Behavior

The torch.einsum operation should execute without any dimension mismatch errors.

Actual Behavior

The torch.einsum operation fails due to a dimension mismatch between def_matrix and vertices.

Debug Information

Here are the shapes of the tensors involved in the torch.einsum operation:

• def_matrix shape: torch.Size([6890, 10475])
• vertices shape: torch.Size([1, 11437, 3])

Relevant Code

Excerpt from def_transfer.py:

def apply_deformation_transfer(
    def_matrix: Tensor,
    vertices: Tensor,
    faces: Tensor,
    use_normals=False
) -> Tensor:
    ''' Applies the deformation transfer on the given meshes
    '''
    if use_normals:
        raise NotImplementedError
    else:
        # Debug prints to check the shapes
        print(f"def_matrix shape: {def_matrix.shape}")
        print(f"vertices shape: {vertices.shape}")
        
        def_vertices = torch.einsum('mn,bni->bmi', [def_matrix, vertices])
        return def_vertices

Additional Context

The issue seems to be caused by a mismatch in the dimensions of def_matrix and vertices. The second dimension of def_matrix (10475) does not match the second dimension of vertices (11437). This discrepancy needs to be resolved to ensure the torch.einsum operation can execute successfully.

I used this config file:

datasets:
    mesh_folder:
        data_folder: 'transfer_data/meshes/smplx'
deformation_transfer_path: 'transfer_data/smplx2smpl_deftrafo_setup.pkl'
mask_ids_fname: ''
summary_steps: 100

edge_fitting:
    per_part: False

optim:
    type: 'lbfgs'
    maxiters: 200
    gtol: 1e-06

body_model:
    model_type: "smpl"
    gender: "female"
    ext: 'pkl'
    folder: "models"
    use_compressed: False
    use_face_contour: True
    smpl:
        betas:
            num: 10

Environment

• OS: Ubuntu
• Python version: 3.8
• PyTorch version: 2.4.1

Labels

• Bug
• RuntimeError
• DimensionMismatch