CUDA out of memory

[notbadzhu]

• When I tried to use cuda to speed up training, an error occurred. Using CPU alone is too slow. How can I solve the problem。
• Environment：Windows 10
• PDF： Probability and Statistics (4th Edition) (Morris H. DeGroot, Mark J. Schervish) (Z-Library).pdf

Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
C:\Users\acer\miniconda3\Lib\site-packages\torch\functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ..\aten\src\ATen\native\TensorShape.cpp:3527.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
Traceback (most recent call last):
  File "D:\github_projects\marker-master\convert_single.py", line 22, in <module>
    full_text, out_meta = convert_single_pdf(fname, model_lst, max_pages=args.max_pages, parallel_factor=args.parallel_factor)
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\github_projects\marker-master\marker\convert.py", line 146, in convert_single_pdf
    filtered, eq_stats = replace_equations(
                         ^^^^^^^^^^^^^^^^^^
  File "D:\github_projects\marker-master\marker\cleaners\equations.py", line 294, in replace_equations
    predictions = get_nougat_text_batched(images, flat_reformat_region_lens, nougat_model, batch_size)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\github_projects\marker-master\marker\cleaners\equations.py", line 109, in get_nougat_text_batched
    model_output = nougat_model.inference(image_tensors=sample, early_stopping=False)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\nougat\model.py", line 580, in inference
    last_hidden_state = self.encoder(image_tensors)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\nougat\model.py", line 123, in forward
    x = self.model.layers(x)
        ^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\torch\nn\modules\container.py", line 215, in forward
    input = module(input)
            ^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\timm\models\swin_transformer.py", line 413, in forward
    x = blk(x)
        ^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\timm\models\swin_transformer.py", line 295, in forward
    attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\acer\miniconda3\Lib\site-packages\timm\models\swin_transformer.py", line 183, in forward
    attn = (q @ k.transpose(-2, -1))
            ~~^~~~~~~~~~~~~~~~~~~~~
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 128.00 MiB. GPU 0 has a total capacty of 4.00 GiB of which 0 bytes is free. Of the allocated memory 3.38 GiB is allocated by PyTorch, and 16.99 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

[VikParuchuri]

This is strange, since marker should work with about 3GB of VRAM. It may be due to the type of card. What you could try is decreasing the batch sizes for various settings:

• LAYOUT_BATCH_SIZE
• ORDERER_BATCH_SIZE
• NOUGAT_BATCH_SIZE

Reducing these will reduce VRAM usage.

You could also try changing the casts to torch.bfloat16 to torch.float16. Basically search torch.bfloat16 and replace with torch.float16. I don't know if earlier cards support bfloat16 (maybe they just use float32 if the card doesn't support bfloat16?).