Unable to finetune BiRefNet_HR-general-epoch_130.pth using public training code.

[hardikjain8]

Hello,

We are trying to fine tune the released HR model checkpoint (BiRefNet_HR-general-epoch_130.pth) using the public repository.

Repository : https;//github.com/ZhengPeng7/BiRefNet

Training command : bash train.sh modelHR 0,1

Our objective is to fine-tune the released HR checkpoint on a custom dataset

When using the BiRefNet_HR-general-epoch_130.pth That training process does not support normally and appears to stall during initialization/loading. We are unable to successfully start fine tuning from released HR checkpoint. Additionally during training we also encountered a issue which is originating from the lost computation. I have attached the screenshot of the issue for your reference.

This suggests that some values passed to the loss function may be outside the expected range although we are unsure whether this is related to

• The HR checkpoint itself
• The current training configuration
• A mismatch between the released HR checkpoint and the public training code.

Any guidance would be greatly appreciated. We would be happy to provide additional logs or debugging information if required.

Thank you.

[ZhengPeng7]

Hi @hardikjain8, what's the GPU you used here?
You need to check the GPU util rate, for example, via nvitop, to monitor it to see if the GPUs are computing there.

The training HR model is slow since the resolution for it is 2048x2048. And the compile at the beginning also takes a lot of time to prepare. You can turn it off in the config for a faster debug.

As shown in the screenshot you attached, you're right -- some values of annotation are out of range [0 ~ 1], which is quite normal to see on custom datasets. I don't think this is a problem with the weights. You can also use existing datasets, which is 100% right for an ablation study. For example, use DIS5K here for training to see if the out-of-range problem disappears.

Let me know if it's fixed after checking your custom datasets.

[hardikjain8]

Hi @ZhengPeng7

Thank you for the suggestions.

I checked the GPU utilization using monitoring tools as recommended The GPU's are allocated and memory is being consumed but the training appears to remain stuck for a very long time before any actual training process is shown.

I also verified the dataset annotations and fix the out of range mask values. After fixing the dataset issue the particularly error no longer appears. However the HR model still seems to hang during startup initialization I understand that the HR model uses 2048x2048 resolution and that compilation may take additional time but the process remains at the same stage for extended without producing training logs or progressing to the first iteration.

I have attached screenshots showing the GPU utilization and memory usage and also the training output where the process appears to stall.

Thank you for your help.

[ZhengPeng7]

As I said before, you can turn off compile for debugging.
The screenshot shows that GPU:0 occupies more memory, which could be caused by the fact that it's still in compile.
Turn it off to see if it starts to run and print logs quickly.

[hardikjain8]

As you recommended I have made self.compile in config.py as "False". The same is shown in second screenshot as well.

[ZhengPeng7]

Oh, it's weird.
You can use a very low resolution and one single GPU to have a try, to eliminate other factors like DDP/Slow steps.

[hardikjain8]

Thank you for your previous guidance.

I would like to provide some additional context about our use case and ask for your recommendation.

Our Use Case

We are building a green-screen background removal system. Our dataset contains adults, children, and strollers/prams for small children.

The main challenges in our application are:

• Preserving fine hair details.
• Maintaining sharp body boundaries.
• Correctly separating foreground regions where multiple people are standing very close together with very little space between them.
• Handling strollers and complex foreground structures without losing details.
• Because of these requirements, we selected BiRefNet-HR since it appeared to provide significantly better boundary quality compared to the other models we evaluated.

Current Training Setup

• GPU: NVIDIA RTX A6000 (48 GB VRAM)
• CPU: 9 vCPUs
• Platform: RunPo
• Number of GPUs: 1
• Checkpoint: "BiRefNet_HR-general-epoch_130.pth"
• Training command: "bash train.sh modelHR 0"
• Checkpoint Loading Issue

However, while attempting to fine-tune using "BiRefNet_HR-general-epoch_130.pth", we are encountering a significant state-dictionary mismatch when we used batch_size = 1.

The checkpoint reports approximately 215 unexpected keys during loading.

Because of this, I am unsure whether:

• The released HR checkpoint is fully compatible with the current training code.
• The checkpoint was trained using a different architecture revision
• Additional HR-specific training code or configuration is required.
• Would you advise us on this because such a large number of unexpected keys (around 215) indicates a deeper incompatibility between the checkpoint and the current architecture.

Memory Constraints

Additionally:

With "batch_size=1", I can proceed to checkpoint loading but encounter the state-dictionary issue.
Increasing to "batch_size=2" results in CUDA out-of-memory errors even on the RTX A6000 (48 GB VRAM).

Memory Optimization

We are currently training with "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True bash train.sh modelHR 0" to reduce CUDA memory fragmentation and make training more memory-efficient on our RTX A6000.

Could you please confirm whether using "expandable_segments:True" has any impact on training quality, convergence, final model performance, or reproducibility? My understanding is that it only affects memory allocation, but I would appreciate your confirmation.

Model Recommendation

We also experimented with Matting-HR, but for our use case we observed:

• Poor handling of strollers/prams
• Dark shadows being classified as foreground.
• Dark objects occasionally leaking into the foreground mask.
• Inferior boundary quality compared to BiRefNet-HR.

Given our use case (green-screen removal, adults, children, strollers, fine hair details, and narrow gaps between people), would you still recommend continuing with BiRefNet-HR, or is there another model that you would suggest evaluating?

Thank you again for your time and support. Any guidance would be greatly appreciated.

[ZhengPeng7]

I've read your detailed replies above.
Batch size (bs) should be >= 2 due to the BN in pretrained models.
Back to the problem, the resolution used in pretraining is not an error for fine-tuning. You can use totally different resolutions between fine-tuning and those used in my previous pretraining.
Therefore, I suspect that the problem you met is caused by the GPU or the system.
To validate in the simplest way -- set a much lower resolution with bs >= 2 to see if training normally prints logs.

For other models, you can also have a try on the BiRefNet_dynamic, which was trained in different resolutions and is robust in those scenarios.

[hardikjain8]

Okay thank you for your response.

To better understand the hardware requirements for BiRefNet-HR training, we tested multiple GPU configurations on RunPod:

These were started for training on 2048 resolution because we need sharpness in our output.

I have attached screenshots from RunPod for reference.

Our use case images also include strollers of this kind

My questions are:

1. Are these memory requirements expected for BiRefNet-HR fine-tuning?
2. Is there a recommended GPU configuration for training/fine-tuning BiRefNet-HR at high resolutions?
3. Why do we face OOM at every GPU?
4. Even if the GPU shows only 100MB needed in A40 which is a 48GB VRAM GPU, it still gets Out of Memory in a GPU of 80GB VRAM?
5. For 2048×2048 fine-tuning, what dataset size would you recommend?
6. What is the minimum number of image-mask pairs required for meaningful improvements?
7. How many epochs are typically recommended for fine-tuning BiRefNet-HR-matting model.

Please provide us the solution for this issue and also your recommended GPU with batch size suitable for 2048 x 2048 resolution.

Inference Comparison

We also performed inference at 2048 resolution using the following checkpoints:

• BiRefNet_HR-general-epoch_130.pth
• BiRefNet_HR-matting-epoch_135.pth
• BiRefNet_dynamic-general-epoch_174.pth
• BiRefNet_dynamic-matting-epoch_159.pth

I have attached the original image together with the outputs from all three checkpoints for comparison.

Original Image :

1. BiRefNet_HR-general-epoch_130.pth

2. BiRefNet_HR-matting-epoch_135.pth

3. BiRefNet_dynamic-general-epoch_174.pth

4. BiRefNet_dynamic-matting-epoch_159.pth

Thank you again for your time and support

[ZhengPeng7]

Well, I used H200s to train the HR versions...
Turning on the compile can decrease the GPU memory needed for training. As I wrote in README, it takes ~24GB memory to train in 1024x1024 and bs==2, which should be around 1/4 of the memory for 2048x2048 + bs==2 training.
For common fine-tuning, 300-500 samples should be enough for acceptable but not very extreme precision. However, HR and matting may need a higher requirement. 30-50 epochs would bring improvement for easy cases, while many epochs are needed for more complicated cases, including HR and matting.

[hardikjain8]

Hi @ZhengPeng7 ,

I wanted to share a quick update and ask a few follow-up questions.

We were able to obtain a NVIDIA RTX PRO 6000 GPU setup with approximately 96 GB VRAM, and training has now started successfully with:

• Resolution: 2048×2048
• Batch Size: 2
• "self.compile = True"
• Checkpoint: BiRefNet_HR-matting-epoch_135.pth

This is the first configuration that has been able to train successfully at 2048 resolution for our use case.

I had a few questions regarding fine-tuning:

1. Our dataset currently contains around 2.5k–3k images. Is this sufficient for fine-tuning, or would you recommend increasing the dataset size further?

2. We are currently planning to train for 50 epochs. Do you think this is reasonable, or would you suggest training for more epochs?

3. Is there a learning rate that you would generally recommend for fine-tuning the HR matting checkpoint on a custom dataset?

4. After every epoch, we save the checkpoint, run evaluation, and append the metrics to a CSV file. What would be the best way to determine which checkpoint is the most suitable for deployment?

Thank you again for your guidance. Your suggestions have been very helpful throughout this process.

[ZhengPeng7]

1. More is better, but keep the high quality of them.
2. As I said, it depends on the difficulty level of the annotation.
3. The default one is okay. If you want more adaptation, decrease it a little bit to see the difference.
4. You definitely need a solid validation set to measure the checkpoints.

You did make me read too much, haha. You're welcome.

[hardikjain8]

Hi @ZhengPeng7,

We have successfully started fine-tuning and would like your advice regarding a dataset-related issue that we are currently observing.

Dataset Setup

For our training dataset, we generated binary masks for all original images. The masks contain only foreground and background regions (0 and 1 values).

Observation

After training with the BiRefNet-HR Matting checkpoint, we noticed that the model struggles on images containing:

• Fuzzy or fine hair details
• Semi-transparent hair boundaries
• Regions where some background is visible through the hair

In these cases, parts of the background tend to remain attached to the foreground prediction.

Comparison with Base Model

Interestingly, when we use the original checkpoint and apply the foreground estimation/refinement technique, the results are noticeably better around hair boundaries.

However, after fine-tuning on our dataset, the output quality around these regions degrades.

Questions

1. For fine-tuning the HR Matting model, would you recommend using soft matte masks instead of strict binary masks, especially for images containing fuzzy hair?

2. Could fine-tuning a matting model on binary masks reduce its ability to predict soft alpha boundaries?

3. For a use case involving people, children, hair details, and green-screen removal, would you recommend:

   • Binary masks
   • Soft alpha mattes

4. Based on the attached examples, do you think the issue is more likely related to the annotations or to some other aspect of training?

I have attached several example images showing the original image, mask, and prediction results for reference.

Thank you again for your guidance.

[ZhengPeng7]

Hi, let me reply to your questions concisely:

1. If you care about the hair regions, you should surely use the setting of the Matting task.
2. No, usually. But it can become worse if the quality of the annotation decreases.
3. When the transparency is important, use alpha mattes. Binary seg is an easier task and easier to train. So, it's often a trade-off.
4. If these gray images are annotations, they need some improvement on the regions with transparency.

[hardikjain8]

Hi @ZhengPeng7

Thank you for your response, all the help has been really useful for us to carry out our project, really appreciate your work.

So currently we are having images of very high size which is around 9MB to 12MB. We want to reduce it nearly to 50-60% for our usecase. If we are doing so using Pngquant, then we are losing the quality of image. So do you suggest any better alternative which could keep up the quality of image and also make it smaller in size around 3MB or 4MB.

Thank you.

[ZhengPeng7]

You are welcome...

But I never did such things and met such a problem. For instance, many images in DIS-TE4 of DIS5K are also > 10MB, but it works well in the project with normal resize functions. So, I'm a bit confused why you need things like PNGQuant.