SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[SegGPT] Fix loss calculation (#30421)

Browse Source 
* Fixed main train issues

* Added loss test

* Update src/transformers/models/seggpt/modeling_seggpt.py

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

* Added missing labels arg in SegGptModel forward

* Fixed typo

* Added slow test to test loss calculation

---------

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>
This commit is contained in:
Eduardo Pacheco
2024-04-24 16:24:34 +02:00
committed by GitHub
parent 37fa1f654f
commit d26c14139c
2 changed files with 89 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
tests/models/seggpt/test_modeling_seggpt.py
View File

@@ -16,6 +16,7 @@
import inspect
import math
import unittest
from datasets import load_dataset
@@ -39,6 +40,7 @@ if is_torch_available():
from torch import nn
from transformers import SegGptForImageSegmentation, SegGptModel
from transformers.models.seggpt.modeling_seggpt import SegGptLoss
if is_vision_available():
@@ -298,6 +300,22 @@ class SegGptModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
model_row_output[key] = model_row_output[key][1:]
recursive_check(model_batched_output[key], model_row_output[key], model_name, key)
def test_seggpt_loss(self):
torch.manual_seed(100)
config = self.model_tester.get_config()
prompt_masks = torch.rand(1, config.num_channels, config.image_size, config.image_size)
label = torch.rand(1, config.num_channels, config.image_size, config.image_size)
pred_masks = torch.rand(1, config.num_channels, config.image_size * 2, config.image_size)
# seq_len x 2 because the loss concatenates prompt_masks and labels as pred_masks is concatenated
bool_masked_pos = torch.rand(1, self.model_tester.seq_length * 2) > 0.5
loss = SegGptLoss(config)
loss_value = loss(prompt_masks, pred_masks, label, bool_masked_pos)
expected_loss_value = torch.tensor(0.3340)
self.assertTrue(torch.allclose(loss_value, expected_loss_value, atol=1e-4))
@slow
def test_model_from_pretrained(self):
model_name = "BAAI/seggpt-vit-large"
@@ -312,6 +330,20 @@ def prepare_img():
return images, masks
def prepare_bool_masked_pos(config: SegGptConfig):
num_patches = math.prod([i // config.patch_size for i in config.image_size])
mask_ratio = 0.75
torch.manual_seed(2)
num_masked_patches = int(num_patches * mask_ratio)
shuffle_idx = torch.randperm(num_patches)
bool_masked_pos = torch.FloatTensor([0] * (num_patches - num_masked_patches) + [1] * num_masked_patches)[
shuffle_idx
]
bool_masked_pos = bool_masked_pos.unsqueeze(0).bool()
return bool_masked_pos
@require_torch
@require_vision
class SegGptModelIntegrationTest(unittest.TestCase):
@@ -390,3 +422,30 @@ class SegGptModelIntegrationTest(unittest.TestCase):
self.assertEqual(outputs.pred_masks.shape, expected_shape)
self.assertTrue(torch.allclose(outputs.pred_masks[0, :, 448:451, :3], expected_slice, atol=4e-4))
@slow
def test_one_shot_with_label(self):
model = SegGptForImageSegmentation.from_pretrained("BAAI/seggpt-vit-large").to(torch_device)
image_processor = self.default_image_processor
images, masks = prepare_img()
input_image = images[1]
label = masks[1]
prompt_image = images[0]
prompt_mask = masks[0]
inputs = image_processor(
images=input_image, prompt_masks=prompt_mask, prompt_images=prompt_image, return_tensors="pt"
).to(torch_device)
labels = image_processor(images=None, prompt_masks=label, return_tensors="pt")["prompt_masks"].to(torch_device)
bool_masked_pos = prepare_bool_masked_pos(model.config).to(torch_device)
with torch.no_grad():
outputs = model(**inputs, labels=labels, bool_masked_pos=bool_masked_pos)
expected_loss = torch.tensor(0.0074).to(torch_device)
self.assertTrue(torch.allclose(outputs.loss, expected_loss, atol=1e-4))