You can deploy this model and start earning money today!

avanishd/
efficient-net-v2-m-finetuned-skin-lesion-classification

Name: efficient-net-v2-m-finetuned-skin-lesion-classification
Availability: InStock

safetensors

model_hub_mixin

pytorch_model_hub_mixin

model-index

Views

No views yet

efficient-net-v2-m-finetuned-skin-lesion-classification

This model is a finetuned version of Pytorch EfficientNet v2-M for skin lesion classification.

Paramas (M): 52.9

Intended Uses & Limitations

Intended Use

This model is intended for dermoscopic skin lesion classification using a 224x224 image size.

Limitations

This model was only trained for 1 epoch and relies on images that are similar in quality to the ISIC dataset. It does not perform well on low-quality images.

How to Get Started with the Model

Python

class EfficientNetSkinLesionClassifier(nn.Module, PyTorchModelHubMixin):
  """
  PytorchModelHubMixin adds push to Hugging Face Hub

  See: https://huggingface.co/docs/hub/models-uploading#upload-a-pytorch-model-using-huggingfacehub
  """
  def __init__(self, num_classes=1):
    super(EfficientNetSkinLesionClassifier, self).__init__()

    # Step 1: Initialize backbone with the best available weights
    base_weights = EfficientNet_V2_M_Weights.DEFAULT
    self.backbone = efficientnet_v2_m(weights=base_weights)

    # Freeze entire backbone
    for param in self.backbone.parameters():
        param.requires_grad = False

    # Unfreeze ONLY the last block (low level feature detection is fine)
    for param in self.backbone.features[-1].parameters():
        param.requires_grad = True

    # Step 3: Replace the original classification head with a new one
    # The original classifier is called self.backbone.classifier
    # See: https://docs.pytorch.org/vision/main/_modules/torchvision/models/efficientnet.html#EfficientNet_V2_M_Weights

    feat_dim = self.backbone.classifier[1].in_features # Access the in_features of the final Linear layer

    # Same classifier format as original
    # Just switiching from ImageNet 1k classes to our binary classification
    # Dropout is 0.3 in pytorch implementation
    self.backbone.classifier = nn.Sequential(
        nn.Dropout(p=0.3, inplace=True),
        nn.Linear(feat_dim, num_classes)
    )

  def forward(self, x):
    out = self.backbone(x)

    return out

from huggingface_hub import hf_hub_download

weights_path = hf_hub_download(
        repo_id="avanishd/efficient-net-v2-m-finetuned-skin-lesion-classification",
        filename="model.safetensors"
    )

from safetensors.torch import load_model

model = EfficientNetSkinLesionClassifier()
load_model(model, filename=weights_path, strict=True)

model.to(device) # Don't forget to put on GPU

model.eval()  # Set model to evaluation mode

# Example with PH2 Dataset

class PH2Dataset(Dataset):
  """
  Dataset for PH2 images, which are in png format.

  PH2 contains skin lesions images classified as

  - Common Nevus (benign)
  - Atypical Nevus (benign)
  - Melanoma (malignant)

  No need for is real label here, since this is purely for testing
  """

  def __init__(self, dir_path, metadata, transform=None):
    super(PH2Dataset, self).__init__()

    self.dir_path = dir_path
    self.transform = transform

    self.image_files = [os.path.join(dir_path, f) for f in os.listdir(dir_path)
                        if f.lower().endswith(('.jpg', '.jpeg', '.png'))]

    # Load metadata w/ polars (only 2 columns)
    self.metadata = pl.read_csv(metadata)

    self.diagnostic_mapping = {
        "Common Nevus": 0,
        "Atypical Nevus": 0,
        "Melanoma": 1,
    }

  def __len__(self):
    return len(self.image_files)

  def __getitem__(self, idx):
    # The image name in the metadata csv are like IMD003
    image_id = self.image_files[idx].split('/')[-1].split('.')[0]

    # Still need the entire path to open the image
    image = Image.open(self.image_files[idx]).convert('RGB')

    if self.transform: # Apply transform if it exists
      image = self.transform(image)

    diagnosis = self.metadata.filter(pl.col("image_name") == image_id).select("diagnosis").item()

    label = torch.tensor(self.diagnostic_mapping[diagnosis], dtype=torch.int16)

    return image, label


transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # Image net mean and std
    transforms.Resize((224, 224)),  # Dimensions for Efficient Net v2
])

ph_2_images = "/content/data/ph2_data/images"
ph_2_metadata = "/content/data/ph2_data/ph_2_dataset.csv"

ex_dataset = PH2Dataset(ph_2_images, ph_2_metadata, transform)

ex_loader = DataLoader(ex_dataset, batch_size=64, shuffle=False)

for (images, labels) in test_loader:
  images = images.to(device)
  labels = labels.to(device)
  output = model(images)

  y_pred_prob = torch.sigmoid(output).cpu().numpy().ravel()
  y_pred = np.where(y_pred_prob < 0.5, 0, 1)

  return y_pred

Training and evaluation data

This model was trained with the ISIC 2024 challenge and ISIC 2024 synthetic datasets.

For the ISIC 2024 Challenge data, an 80-20 train test split was applied, and the test split was used to evaluate the model.

Training Procedure

Training hyperparameters

learning_rate: 0.001
train_batch_size: 64
eval_batch_size: 64
optimizer: Use OptimizerNames.ADAM
num_epochs: 1

Training results

Training Loss	Epoch	Step
0.4121	1	100
0.2303	1	200
0.7978	1	300
0.5703	1	400
0.3542	1	500
0.5727	1	600
0.4968	1	700
0.2870	1	800
0.5946	1	900
0.1645	1	1000
0.3951	1	1100
0.4131	1	1200
0.1167	1	1300
0.1780	1	1400
0.3001	1	1500
0.5137	1	1600
0.1442	1	1700
0.2861	1	1800
1.1042	1	1900
0.6788	1	2000
0.3203	1	2100
0.3083	1	2200
0.6786	1	2300
0.7330	1	2400
3.2796	1	2500
0.1256	1	2600
0.5112	1	2700
0.5217	1	2800
0.2267	1	2900
0.2648	1	3000
0.4319	1	3100
0.1849	1	3200
0.1746	1	3300
0.2552	1	3400
0.1613	1	3500
0.4685	1	3600
0.1714	1	3700
0.2739	1	3800
0.2142	1	3900
0.2366	1	4000
0.1659	1	4100
0.1367	1	4200
0.3813	1	4300
0.3124	1	4400
0.3023	1	4500
0.2140	1	4600
0.9778	1	4700
0.3501	1	4800
0.2174	1	4900
0.3568	1	5000
0.0863	1	5100
6.8996	1	5200

Framework versions

Pytorch 2.9.0+cu126
torchvision: 0.24.0+cu126
numpy: 2.0.2