Fine tune TSR - Githubissues

Hi, I've been testing your code in order to fine tune my own version of triposr with a custom dataset. However, as I try to set the model weights in serv.py in this code segment (lines 195 to 199), by uncommenting some lines:

    model = TSR(img_size=image_size, depth=16//2, embed_dim=768, num_channels=1024, num_layers=16//2, cross_attention_dim=768, radius=99, valid_thresh=0.00001, num_samples_per_ray=128, n_hidden_layers=9, official=False)
    #model.load_state_dict(torch.load('./ckpt/TripoSR/model.ckpt', map_location='cpu'))
    model.to(device)
    model.train()

I get the following error:

/data/2024-I/MACHINE_LEARNING_TECHNIQUES/project_v2/triposr/tripovenv/lib/python3.11/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  _torch_pytree._register_pytree_node(
/data/2024-I/MACHINE_LEARNING_TECHNIQUES/project_v2/triposr/tripovenv/lib/python3.11/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  _torch_pytree._register_pytree_node(
parameters: 399 M
1
0it [00:00, ?it/s]
Traceback (most recent call last):
  File "/data/2024-I/MACHINE_LEARNING_TECHNIQUES/project_v2/triposr/super-ai-vision-stereo-world-generate-triposr/superv.py", line 275, in <module>
    main()
  File "/data/2024-I/MACHINE_LEARNING_TECHNIQUES/project_v2/triposr/super-ai-vision-stereo-world-generate-triposr/superv.py", line 272, in main
    train(image_size=[512,64][1], batch_size=1, epochs=1, checkpoint_path='./outs/ckpt/', best_checkpoint_file='./outs/ckpt/checkpoint.pth', device=device)
  File "/data/2024-I/MACHINE_LEARNING_TECHNIQUES/project_v2/triposr/super-ai-vision-stereo-world-generate-triposr/superv.py", line 243, in train
    image, alpha = model.renderer(model.decoder, scene_code, rays_o, rays_d)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/2024-I/MACHINE_LEARNING_TECHNIQUES/project_v2/triposr/tripovenv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/2024-I/MACHINE_LEARNING_TECHNIQUES/project_v2/triposr/tripovenv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/2024-I/MACHINE_LEARNING_TECHNIQUES/project_v2/triposr/super-ai-vision-stereo-world-generate-triposr/network_nerf_renderer.py", line 190, in forward
    xyz = (rays_o[:, None, :] + z_vals[..., None] * rays_d[..., None, :])  # (N_rays, N_sample, 3)
                                ~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (0) must match the size of tensor b (4096) at non-singleton dimension 0

The training seems to work well with the lines commented as in the original file in the repository, in a 2080TI, and this is the only training code I've been able to run to generate 3D models from images. I would like to know if you have had a similar problem. Also, I had to download the file directly from huggingfaces since the wget command did not work.

Thanks in advance, Johan

yuedajiong / super-ai-vision-stereo-world-generate-triposr

Fine tune TSR #1