Closed hihxy closed 4 months ago
hello @hihxy , thanks for the question.
Have you tried modifying the --batch-size
argument of inference_grid.py
? If you did and it wasn't helpful, let us know and we'll think of ways to make the script more efficient / parallelizable.
Also just as a sidenote -- the reason why the training runs OOM at 5K is because we start applying the contrastive and cycle-consistency losses in the refined feature space after 5K, which demands more memory.
Thanks for your serious reply @tnarek .
In the training step, I only changed the train_batch_size
parameter to 128 in the ./config/train.yaml
file, and it worked.
########################## data loader ##########################
video_resw: 854
video_resh: 476
fg_traj_ratio: 0.5
keep_traj_in_cpu: true # set to true for long videos (> 150 frames)
train_batch_size: 128 # defult 512
batch_n_frames: 4
# sampler_batch_iterations: 1000 # uncomment for large videos
And in the inference step, the ./config/train.yaml
file in trainning step ( train_batch_size: 128 / 64 /32 I have tried ). I think it will change the --batch-size
argument of inference_grid.py
. Is it right?
After reading your reply, I have used the code above to infer, and changed the --batch-size
(128,64,32...) but still runs OOM.
python ./inference_grid.py \
--config ./config/train.yaml \
--data-path ./dataset/ljc \
--batch-size 64 \
--use-segm-mask \
> ./logs/inference.log 2>&1
The inference log output is the same as my first issue.
Traceback (most recent call last):
File "/data/hxy/dino-tracker/./inference_grid.py", line 54, in <module>
run(args)
File "/data/hxy/envs/dino-tracker/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/data/hxy/dino-tracker/./inference_grid.py", line 25, in run
model_inference = ModelInference(
File "/data/hxy/dino-tracker/models/model_inference.py", line 91, in __init__
self.model.cache_refined_embeddings()
File "/data/hxy/dino-tracker/models/tracker.py", line 132, in cache_refined_embeddings
refined_features, _ = self.get_refined_embeddings(torch.arange(0, self.video.shape[0]))
File "/data/hxy/dino-tracker/models/tracker.py", line 125, in get_refined_embeddings
refined_embeddings = frames_dino_embeddings + residual_embeddings
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.31 GiB. GPU 0 has a total capacty of 23.69 GiB of which 2.34 GiB is free. Including non-PyTorch memory, this process has 21.21 GiB memory in use. Of the allocated memory 20.89 GiB is allocated by PyTorch, and 23.22 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
now I see that the reason for OOM is the caching of the refined embeddings. We do this to compute them only once for the entire video, and reuse them for all query points.
You can turn this line off in model_inference.py
(line 91), but it will slow down the inference.
Alternatively, trajectories can be computed for a subset of 100 frames separately, which will demand caching less refined embeddings.
Thank you for your reply. I have tried to turn off this line, but it will also lead to OOM problem.
model_inference.py
class ModelInference(torch.nn.Module):
def __init__(
self,
model: Tracker,
range_normalizer: RangeNormalizer,
anchor_cosine_similarity_threshold: float = 0.5,
cosine_similarity_threshold: float = 0.5,
) -> None:
super().__init__()
self.model = model
self.model.eval()
# self.model.cache_refined_embeddings()
self.range_normalizer = range_normalizer
self.anchor_cosine_similarity_threshold = anchor_cosine_similarity_threshold
self.cosine_similarity_threshold = cosine_similarity_threshold
log information
nohup: ignoring input
/data/hxy/envs/dino-tracker/lib/python3.9/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:3526.)
return _VF.meshgrid(tensors, **kwargs) # type: ignore[attr-defined]
Traceback (most recent call last):
File "/data/hxy/dino-tracker/./inference_grid.py", line 54, in <module>
run(args)
File "/data/hxy/envs/dino-tracker/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/data/hxy/dino-tracker/./inference_grid.py", line 39, in run
grid_trajectories, grid_occlusions = model_inference.infer(grid_query_points, batch_size=args.batch_size)
File "/data/hxy/envs/dino-tracker/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/data/hxy/dino-tracker/models/model_inference.py", line 212, in infer
trajs = self.compute_trajectories(query_points, batch_size) # N x T x 3
File "/data/hxy/dino-tracker/models/model_inference.py", line 98, in compute_trajectories
trajecroies = generate_trajectories(
File "/data/hxy/envs/dino-tracker/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/data/hxy/dino-tracker/models/model_inference.py", line 70, in generate_trajectories
trajectory_pred = generate_trajectory(query_point=query_point, video=video, model=model, range_normalizer=range_normalizer, dst_range=dst_range, use_raw_features=use_raw_features,
File "/data/hxy/envs/dino-tracker/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/data/hxy/dino-tracker/models/model_inference.py", line 51, in generate_trajectory
trajectory_coordinate_preds_normalized = model(trajectory_input, use_raw_features=use_raw_features)
File "/data/hxy/envs/dino-tracker/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/data/hxy/envs/dino-tracker/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "/data/hxy/dino-tracker/models/tracker.py", line 319, in forward
frame_embeddings, residual_embeddings, raw_embeddings = self.get_refined_embeddings(frames_set_t, return_raw_embeddings=True)
File "/data/hxy/dino-tracker/models/tracker.py", line 120, in get_refined_embeddings
residual_embeddings = torch.zeros_like(frames_dino_embeddings)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.77 GiB. GPU 0 has a total capacty of 23.69 GiB of which 3.21 GiB is free. Including non-PyTorch memory, this process has 20.35 GiB memory in use. Of the allocated memory 20.02 GiB is allocated by PyTorch, and 12.06 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
May I ask is it because the tensor of frames_dino_embeddings
is too large? How should I optimize it? By the way, does the dino-tracker support training and interfering with multiple GPUs? Look forward to your reply, thank you so mush.
you are correct, the problem is that the frame embeddings do not fit in the GPU memory. The OOM happens even before computing the embeddings in line 120 of tracker.py
, where we initialize a tensor of zeros where the residual embeddings are stored later on. So if this tensor does not fit in memory, the frame embeddings will not fit either.
currently, our code does not support parallel GPU's. a simple solution can be to store the raw DINO embeddings to a second GPU device, and move them to the first device only when they are necessary. You can see that before the OOM happens, the GPU already has 20 GB occupied, most of which I suspect are the pre-computed raw DINO embeddings.
Thanks.
Thanks for proposing such a great method! I've been trying to replicate your code recently, but I've run into some issues. CUDA out of memory error occurred during the model inference step. 200 images with a video resolution of 854x476.The GPU used was an NVIDIA RTX3090 with a video memory size of 24GB. The issue also occurred during the training step.
I managed to fix it by changing the
train_batch_size
parameter to 128 in the./config/train.yaml
file, and that sorted things out, wrapping up the training successfully. But then, when I tried using this trained model for inference, the issue resurfaced, and tweaking the parameters didn't help either.So, how can I deal with it? I really need your help.