# How to train a Semantic Segmentation model using transformers tensorflow2 API

**URL:** <https://discuss.huggingface.co/t/how-to-train-a-semantic-segmentation-model-using-transformers-tensorflow2-api/24317>\
**Category:** 🤗Transformers\
**Created:** [October 12, 2022, 8:18am UTC](https://discuss.huggingface.co/t/how-to-train-a-semantic-segmentation-model-using-transformers-tensorflow2-api/24317 "2022-10-12T08:18:26Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![JackieLee](https://avatars.discourse-cdn.com/v4/letter/j/bbe5ce/32.png) [@JackieLee](https://discuss.huggingface.co/u/JackieLee)\
**Post date:** [October 12, 2022, 8:18am UTC](https://discuss.huggingface.co/t/how-to-train-a-semantic-segmentation-model-using-transformers-tensorflow2-api/24317/1 "2022-10-12T08:18:26Z")

</div>

I want to train a Semantic Segmentation model using transformers tensorflow2 API. But the docs only provides the case for pytorch. I tried to implement, but it didn’t work. Here is my code.

```auto
from datasets import load_dataset
from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation

ds = load_dataset("scene_parse_150", split="train[:50]")
ds = ds.train_test_split(test_size=0.2)
train_ds = ds["train"]
test_ds = ds["test"]

import json
from huggingface_hub import cached_download, hf_hub_url

repo_id = "huggingface/label-files"
filename = "ade20k-hf-doc-builder.json"
id2label = json.load(open(cached_download(hf_hub_url(repo_id, filename, repo_type="dataset")), "r"))
id2label = {int(k): v for k, v in id2label.items()}
label2id = {v: k for k, v in id2label.items()}
num_labels = len(id2label)

def data_transforms(example_batch):
    images = [x for x in example_batch["image"]]
    labels = [x for x in example_batch["annotation"]]
    inputs = feature_extractor(images, labels)
    return inputs

train_ds.set_transform(data_transforms)
test_ds.set_transform(data_transforms)

feature_extractor = AutoFeatureExtractor.from_pretrained("nvidia/mit-b0", reduce_labels=True)
model = TFSegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")

from transformers import create_optimizer
optimizer, schedule = create_optimizer(init_lr=0.0001, num_warmup_steps=100,
                                               num_train_steps=100)

model.compile(optimizer=optimizer)

model.fit(x=trainset, validation_data=valset, epochs=1)

```