Initial commit
This commit is contained in:
27
.gitignore
vendored
Normal file
27
.gitignore
vendored
Normal file
@@ -0,0 +1,27 @@
|
|||||||
|
|
||||||
|
# Python cache
|
||||||
|
__pycache__/
|
||||||
|
*.pyc
|
||||||
|
*.pyo
|
||||||
|
|
||||||
|
# Virtual environment
|
||||||
|
venv/
|
||||||
|
.venv/
|
||||||
|
bin/
|
||||||
|
lib/
|
||||||
|
lib64/
|
||||||
|
include/
|
||||||
|
share/
|
||||||
|
pyvenv.cfg
|
||||||
|
|
||||||
|
# Data & datasets
|
||||||
|
data/
|
||||||
|
bin/
|
||||||
|
|
||||||
|
# Model weights
|
||||||
|
*.pth
|
||||||
|
|
||||||
|
# System / logs
|
||||||
|
.DS_Store
|
||||||
|
*.log
|
||||||
|
*.tmp
|
||||||
81
Data.py
Normal file
81
Data.py
Normal file
@@ -0,0 +1,81 @@
|
|||||||
|
from torchvision import datasets, transforms, models
|
||||||
|
import torch
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
# transform images to size
|
||||||
|
def transform(res):
|
||||||
|
return transforms.Compose([
|
||||||
|
# ResNet expects 224 x 224 res
|
||||||
|
transforms.Resize((res, res)),
|
||||||
|
transforms.RandomHorizontalFlip(),
|
||||||
|
transforms.ToTensor(),
|
||||||
|
# normalise to
|
||||||
|
transforms.Normalize(
|
||||||
|
mean=[0.485, 0.456, 0.406],
|
||||||
|
std=[0.229, 0.224, 0.225]
|
||||||
|
)
|
||||||
|
])
|
||||||
|
|
||||||
|
# Load data with 'identity' as target and transform it
|
||||||
|
def get_set(res):
|
||||||
|
return datasets.CelebA(
|
||||||
|
root='./data',
|
||||||
|
split='all',
|
||||||
|
target_type='identity',
|
||||||
|
download=True,
|
||||||
|
transform=transform(res)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def get_ids_and_counts(dataset):
|
||||||
|
return torch.unique(
|
||||||
|
dataset.identity,
|
||||||
|
return_counts=True
|
||||||
|
)
|
||||||
|
|
||||||
|
# filter selected identities from dataset
|
||||||
|
# How many classes, how many images per class
|
||||||
|
def select_ids( dataset, sample_size, class_size):
|
||||||
|
ids, counts = get_ids_and_counts(dataset=dataset)
|
||||||
|
eligible_mask = counts >= sample_size
|
||||||
|
eligible_ids = ids[eligible_mask].numpy()
|
||||||
|
|
||||||
|
if len(eligible_ids) < class_size:
|
||||||
|
raise ValueError(
|
||||||
|
f"Only found {len(eligible_ids)} identities with {sample_size}+ images."
|
||||||
|
)
|
||||||
|
|
||||||
|
# Randomly select 50 identities
|
||||||
|
return np.random.choice(eligible_ids, class_size, replace=False)
|
||||||
|
|
||||||
|
# optional function to get max amount of samples per class
|
||||||
|
def select_balanced_ids(dataset, class_size):
|
||||||
|
ids, counts = get_ids_and_counts(dataset=dataset)
|
||||||
|
|
||||||
|
# sort by number of images (descending)
|
||||||
|
sorted_indices = torch.argsort(counts, descending=True)
|
||||||
|
|
||||||
|
top_ids = ids[sorted_indices][:class_size].numpy()
|
||||||
|
|
||||||
|
return np.array(top_ids, dtype=int)
|
||||||
|
|
||||||
|
|
||||||
|
# split class images to train and test set.
|
||||||
|
def get_indices(dataset, identities, split_at):
|
||||||
|
train_indices = []
|
||||||
|
test_indices = []
|
||||||
|
|
||||||
|
#training_sample = int(sample_size * training_ratio)
|
||||||
|
|
||||||
|
for person_id in identities:
|
||||||
|
# Get all indices for this specific person
|
||||||
|
indices = torch.where(dataset.identity == person_id)[0].numpy()
|
||||||
|
|
||||||
|
# Shuffle the indices for this person
|
||||||
|
np.random.shuffle(indices)
|
||||||
|
|
||||||
|
# split data to testing and training
|
||||||
|
train_indices.extend(indices[:split_at])
|
||||||
|
test_indices.extend(indices[split_at:])
|
||||||
|
|
||||||
|
return train_indices, test_indices
|
||||||
13
IdentitySubset.py
Normal file
13
IdentitySubset.py
Normal file
@@ -0,0 +1,13 @@
|
|||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
class IdentitySubset(torch.utils.data.Dataset):
|
||||||
|
def __init__(self, full_ds, indices, id_mapping):
|
||||||
|
self.full_ds = full_ds
|
||||||
|
self.indices = indices
|
||||||
|
self.id_mapping = id_mapping
|
||||||
|
def __getitem__(self, idx):
|
||||||
|
img, old_id = self.full_ds[self.indices[idx]]
|
||||||
|
return img, self.id_mapping[old_id.item()]
|
||||||
|
def __len__(self):
|
||||||
|
return len(self.indices)
|
||||||
67
Predict.py
Normal file
67
Predict.py
Normal file
@@ -0,0 +1,67 @@
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
@torch.inference_mode() # More memory-efficient than no_grad()
|
||||||
|
def get_loss_per_sample(model, data_loader, device):
|
||||||
|
"""
|
||||||
|
Returns a list of individual losses for every sample in the loader.
|
||||||
|
Useful for MIA to see how 'certain' the model is about specific images.
|
||||||
|
"""
|
||||||
|
model.eval()
|
||||||
|
criterion = torch.nn.CrossEntropyLoss(reduction='none') # Crucial: returns loss per image
|
||||||
|
all_losses = []
|
||||||
|
|
||||||
|
for inputs, labels in data_loader:
|
||||||
|
inputs, labels = inputs.to(device), labels.to(device)
|
||||||
|
|
||||||
|
outputs = model(inputs)
|
||||||
|
|
||||||
|
# Calculate loss for each image in the batch individually
|
||||||
|
loss = criterion(outputs, labels)
|
||||||
|
|
||||||
|
all_losses.extend(loss.cpu().numpy())
|
||||||
|
|
||||||
|
return all_losses
|
||||||
|
|
||||||
|
|
||||||
|
@torch.inference_mode()
|
||||||
|
def get_losses_by_class(model, data_loader, device):
|
||||||
|
"""
|
||||||
|
Returns a dictionary: { class_id: [list_of_losses_for_this_class] }
|
||||||
|
"""
|
||||||
|
model.eval()
|
||||||
|
criterion = torch.nn.CrossEntropyLoss(reduction='none')
|
||||||
|
|
||||||
|
class_losses = {}
|
||||||
|
|
||||||
|
for inputs, labels in data_loader:
|
||||||
|
inputs, labels = inputs.to(device), labels.to(device)
|
||||||
|
outputs = model(inputs)
|
||||||
|
|
||||||
|
# Get individual losses
|
||||||
|
losses = criterion(outputs, labels).cpu().numpy()
|
||||||
|
labels_np = labels.cpu().numpy()
|
||||||
|
|
||||||
|
for i, class_id in enumerate(labels_np):
|
||||||
|
if class_id not in class_losses:
|
||||||
|
class_losses[class_id] = []
|
||||||
|
class_losses[class_id].append(losses[i])
|
||||||
|
|
||||||
|
return class_losses
|
||||||
|
|
||||||
|
|
||||||
|
# evaluate MIA
|
||||||
|
def eval_MIA(forgotten_losses, never_seen_losses):
|
||||||
|
avg_f_loss = np.mean(forgotten_losses)
|
||||||
|
avg_ns_loss = np.mean(never_seen_losses)
|
||||||
|
|
||||||
|
print(f"Average Loss on Forgotten Identity: {avg_f_loss:.4f}")
|
||||||
|
print(f"Average Loss on Unknown Identities: {avg_ns_loss:.4f}")
|
||||||
|
|
||||||
|
if avg_f_loss < avg_ns_loss * 0.8:
|
||||||
|
print("MIA Warning: Model still shows high certainty on forgotten data.")
|
||||||
|
else:
|
||||||
|
print("MIA Success: Model treats forgotten data like unknown data.")
|
||||||
42
ReadME.md
Normal file
42
ReadME.md
Normal file
@@ -0,0 +1,42 @@
|
|||||||
|
# Python venv
|
||||||
|
Start a python environment here in this directory
|
||||||
|
```py
|
||||||
|
python -m venv .
|
||||||
|
```
|
||||||
|
|
||||||
|
Then we start the env using
|
||||||
|
```py
|
||||||
|
source ./bin/activate
|
||||||
|
```
|
||||||
|
|
||||||
|
We can then install whats needed with `pip`. for exampe
|
||||||
|
we can put all dependencies in some text file. say dependencies.txt
|
||||||
|
```py
|
||||||
|
# pip install
|
||||||
|
# already added dependencies.txt
|
||||||
|
pip install -r dependencies.txt
|
||||||
|
|
||||||
|
```
|
||||||
|
|
||||||
|
Downloading the data from google drive was impossible. So Downloaded them manualy
|
||||||
|
and They need to be put in the a ./data directory
|
||||||
|
The download url was available in the error log.
|
||||||
|
`https://drive.google.com/uc?id=0B7EVK8r0v71pZjFTYXZWM3FlRnM`
|
||||||
|
this is the same location thats available in the official site
|
||||||
|
|
||||||
|
```
|
||||||
|
|
||||||
|
```
|
||||||
|
Root_dir/
|
||||||
|
└── data/
|
||||||
|
└── celeba/
|
||||||
|
├── img_align_celeba.zip
|
||||||
|
├── list_attr_celeba.txt
|
||||||
|
├── list_bbox_celeba.txt
|
||||||
|
├── list_eval_partition.txt
|
||||||
|
└── list_landmarks_align_celeba.txt
|
||||||
|
|
||||||
|
```
|
||||||
|
|
||||||
|
once this is done manually
|
||||||
|
|
||||||
14
SetUp.py
Normal file
14
SetUp.py
Normal file
@@ -0,0 +1,14 @@
|
|||||||
|
##
|
||||||
|
import torch
|
||||||
|
from torchvision import datasets, transforms, models
|
||||||
|
|
||||||
|
def get_device():
|
||||||
|
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
# clear cach to boost memory
|
||||||
|
# for new round
|
||||||
|
torch.cuda.empty_cache()
|
||||||
|
return torch.device("cuda")
|
||||||
|
else:
|
||||||
|
return torch.device("cpu")
|
||||||
|
|
||||||
112
Tune.py
Normal file
112
Tune.py
Normal file
@@ -0,0 +1,112 @@
|
|||||||
|
import torch
|
||||||
|
from torch.utils.data import DataLoader
|
||||||
|
from sklearn.metrics import classification_report
|
||||||
|
import SetUp
|
||||||
|
from Data import *
|
||||||
|
from IdentitySubset import IdentitySubset
|
||||||
|
# models
|
||||||
|
from architectures.Model import Model, Architecture
|
||||||
|
|
||||||
|
# numbre of classes
|
||||||
|
CLASS_SIZE = 30
|
||||||
|
# batch
|
||||||
|
BATCH_SIZE = 16
|
||||||
|
|
||||||
|
# size of images per class trainset + testset
|
||||||
|
# 30 works best, more than that and we dont have enough data
|
||||||
|
SAMPLE_SIZE = 30
|
||||||
|
|
||||||
|
# this is then full sample - test sample
|
||||||
|
TRAINING_SMPLE = 28
|
||||||
|
|
||||||
|
# learning rate
|
||||||
|
LR_RATE = 0.0001
|
||||||
|
EPOCHS = 20
|
||||||
|
|
||||||
|
# depends on model architecture
|
||||||
|
# ResNet, DenseNet = 224
|
||||||
|
# Inception = 299
|
||||||
|
RESOLUTION = 299
|
||||||
|
|
||||||
|
# model architecture
|
||||||
|
arch = Architecture.INCEPTION
|
||||||
|
|
||||||
|
# DATA PREPARATION
|
||||||
|
# load data set and prepare
|
||||||
|
dataset = get_set(res = RESOLUTION)
|
||||||
|
# select identities for experiment
|
||||||
|
selected_identities = select_ids(
|
||||||
|
dataset = dataset,
|
||||||
|
sample_size = SAMPLE_SIZE,
|
||||||
|
class_size = CLASS_SIZE
|
||||||
|
)
|
||||||
|
|
||||||
|
print(f'> Selected {CLASS_SIZE} random identity classes from CelebA dataset.')
|
||||||
|
print(f'> A class has {TRAINING_SMPLE} train and {SAMPLE_SIZE-TRAINING_SMPLE} test samples')
|
||||||
|
|
||||||
|
# split class images to train/test indices
|
||||||
|
train_indices, test_indices = get_indices(
|
||||||
|
dataset = dataset,
|
||||||
|
identities = selected_identities,
|
||||||
|
split_at = TRAINING_SMPLE
|
||||||
|
)
|
||||||
|
|
||||||
|
# helps map class id to index
|
||||||
|
id_map = {old_id: new_id for new_id, old_id in enumerate(selected_identities)}
|
||||||
|
|
||||||
|
# we remap identities because crossEntropyLoss requires in indices 0 -> (n-1)
|
||||||
|
# where n = class size.
|
||||||
|
train_data = IdentitySubset(
|
||||||
|
dataset,
|
||||||
|
train_indices,
|
||||||
|
id_map)
|
||||||
|
|
||||||
|
train_loader = DataLoader(
|
||||||
|
train_data,
|
||||||
|
batch_size = BATCH_SIZE,
|
||||||
|
shuffle = True)
|
||||||
|
|
||||||
|
print(f"> Total training images: {len(train_data)}")
|
||||||
|
|
||||||
|
print(f'> Constants : Classes = {CLASS_SIZE}, Batch = {BATCH_SIZE}, epochs = {EPOCHS}')
|
||||||
|
|
||||||
|
# MODEL PREPARATION
|
||||||
|
# cuda if exists (it does here)
|
||||||
|
device = SetUp.get_device()
|
||||||
|
# Create model using Factory
|
||||||
|
model = Model.create(
|
||||||
|
arch = arch,
|
||||||
|
device = device,
|
||||||
|
size = CLASS_SIZE)
|
||||||
|
|
||||||
|
# FINETUNING
|
||||||
|
model.train(
|
||||||
|
epochs = EPOCHS,
|
||||||
|
loader = train_loader,
|
||||||
|
rate = LR_RATE)
|
||||||
|
|
||||||
|
# save.
|
||||||
|
torch.save(
|
||||||
|
model.get().state_dict(),
|
||||||
|
f'{arch.name}.pth')
|
||||||
|
|
||||||
|
# done tuning
|
||||||
|
print('Model saved!')
|
||||||
|
|
||||||
|
# EVALUATE
|
||||||
|
# Testing
|
||||||
|
test_data = IdentitySubset(
|
||||||
|
dataset,
|
||||||
|
test_indices,
|
||||||
|
id_map)
|
||||||
|
|
||||||
|
test_loader = DataLoader(
|
||||||
|
test_data,
|
||||||
|
batch_size=BATCH_SIZE,
|
||||||
|
shuffle=False)
|
||||||
|
|
||||||
|
print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}")
|
||||||
|
|
||||||
|
# Evaluate
|
||||||
|
model.evaluate(
|
||||||
|
loader = test_loader)
|
||||||
15
architectures/DenseNet121.py
Normal file
15
architectures/DenseNet121.py
Normal file
@@ -0,0 +1,15 @@
|
|||||||
|
|
||||||
|
import torch.nn as nn
|
||||||
|
from torchvision import models
|
||||||
|
from architectures.Model import Model
|
||||||
|
|
||||||
|
class DenseNet121(Model):
|
||||||
|
def get(self):
|
||||||
|
|
||||||
|
# load pretrained
|
||||||
|
m = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT)
|
||||||
|
# will modify only the final layers
|
||||||
|
num_ftrs = m.classifier.in_features
|
||||||
|
m.classifier = nn.Linear(num_ftrs, self.size)
|
||||||
|
|
||||||
|
return m.to(self.device)
|
||||||
47
architectures/Inception.py
Normal file
47
architectures/Inception.py
Normal file
@@ -0,0 +1,47 @@
|
|||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.optim as optim
|
||||||
|
from torchvision import models
|
||||||
|
import time
|
||||||
|
|
||||||
|
# Base model
|
||||||
|
from architectures.Model import Model
|
||||||
|
|
||||||
|
class Inception(Model):
|
||||||
|
def get(self):
|
||||||
|
model = models.inception_v3(weights=models.Inception_V3_Weights.DEFAULT)
|
||||||
|
#for param in model.parameters():
|
||||||
|
# param.requires_grad = False
|
||||||
|
model.fc = nn.Linear(model.fc.in_features, self.size)
|
||||||
|
model.AuxLogits.fc = nn.Linear(model.AuxLogits.fc.in_features, self.size)
|
||||||
|
return model.to(self.device)
|
||||||
|
|
||||||
|
def train(self, epochs, loader, rate):
|
||||||
|
# Override because Inception returns a tuple (main, aux)
|
||||||
|
criterion = nn.CrossEntropyLoss()
|
||||||
|
optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=rate)
|
||||||
|
|
||||||
|
print(f"Starting training on {self.device}...")
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
self.model.train()
|
||||||
|
|
||||||
|
for epoch in range(epochs):
|
||||||
|
total_loss = 0.0
|
||||||
|
for inputs, labels in loader:
|
||||||
|
|
||||||
|
inputs, labels = inputs.to(self.device), labels.to(self.device)
|
||||||
|
optimizer.zero_grad()
|
||||||
|
|
||||||
|
outputs, aux_outputs = self.model(inputs)
|
||||||
|
loss = criterion(outputs, labels) + 0.3 * criterion(aux_outputs, labels)
|
||||||
|
|
||||||
|
loss.backward()
|
||||||
|
optimizer.step()
|
||||||
|
total_loss += loss.item()
|
||||||
|
|
||||||
|
print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(loader):.4f}")
|
||||||
|
|
||||||
|
if self.device.type == 'cuda': torch.cuda.synchronize()
|
||||||
|
print(f"Training completed in: {time.time() - start_time:.2f}s")
|
||||||
99
architectures/Model.py
Normal file
99
architectures/Model.py
Normal file
@@ -0,0 +1,99 @@
|
|||||||
|
from abc import ABC, abstractmethod
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.optim as optim
|
||||||
|
import time
|
||||||
|
import numpy as np
|
||||||
|
from sklearn.metrics import classification_report
|
||||||
|
|
||||||
|
class Model(ABC):
|
||||||
|
def __init__(self, device, size):
|
||||||
|
self.device = device
|
||||||
|
self.size = size
|
||||||
|
self.model = self.get()
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
def get(self):
|
||||||
|
# return the model
|
||||||
|
return self.model
|
||||||
|
|
||||||
|
def train(self, epochs, loader, rate):
|
||||||
|
criterion = nn.CrossEntropyLoss()
|
||||||
|
optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=rate)
|
||||||
|
|
||||||
|
print(f"Starting training on {self.device}...")
|
||||||
|
start_time = time.time()
|
||||||
|
self.model.train()
|
||||||
|
|
||||||
|
for epoch in range(epochs):
|
||||||
|
total_loss = 0.0
|
||||||
|
for inputs, labels in loader:
|
||||||
|
inputs, labels = inputs.to(self.device), labels.to(self.device)
|
||||||
|
optimizer.zero_grad()
|
||||||
|
outputs = self.model(inputs)
|
||||||
|
loss = criterion(outputs, labels)
|
||||||
|
loss.backward()
|
||||||
|
optimizer.step()
|
||||||
|
total_loss += loss.item()
|
||||||
|
|
||||||
|
print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss / len(loader):.4f}")
|
||||||
|
|
||||||
|
if self.device.type == 'cuda': torch.cuda.synchronize()
|
||||||
|
print(f"Training completed in: {time.time() - start_time:.2f}s")
|
||||||
|
|
||||||
|
def evaluate(self, loader):
|
||||||
|
self.model.eval()
|
||||||
|
all_preds, all_labels = [], []
|
||||||
|
print("\nEvaluating...")
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
for inputs, labels in loader:
|
||||||
|
inputs, labels = inputs.to(self.device), labels.to(self.device)
|
||||||
|
outputs = self.model(inputs)
|
||||||
|
_, predicted = torch.max(outputs, 1)
|
||||||
|
all_preds.extend(predicted.cpu().numpy())
|
||||||
|
all_labels.extend(labels.cpu().numpy())
|
||||||
|
|
||||||
|
accuracy = 100 * (np.array(all_preds) == np.array(all_labels)).sum() / len(all_labels)
|
||||||
|
print(f"Test Accuracy: {accuracy:.2f}%")
|
||||||
|
print(classification_report(all_labels, all_preds, zero_division=0))
|
||||||
|
|
||||||
|
|
||||||
|
# Using the factory patern here
|
||||||
|
@staticmethod
|
||||||
|
def create(arch, device, size):
|
||||||
|
print(f'>> MODEL ARCHITECTURE >> {arch.name}.')
|
||||||
|
|
||||||
|
match arch:
|
||||||
|
|
||||||
|
# ResNet18
|
||||||
|
case Architecture.RESNET18:
|
||||||
|
from architectures.ResNet18 import ResNet18
|
||||||
|
return ResNet18(device, size)
|
||||||
|
|
||||||
|
# ResNet50
|
||||||
|
case Architecture.RESNET50:
|
||||||
|
from architectures.ResNet18 import ResNet18
|
||||||
|
return ResNet18(device, size)
|
||||||
|
|
||||||
|
# INCEPTION
|
||||||
|
case Architecture.INCEPTION:
|
||||||
|
from architectures.Inception import Inception
|
||||||
|
return Inception(device, size)
|
||||||
|
|
||||||
|
# DENSENET121
|
||||||
|
case Architecture.DENSENET121:
|
||||||
|
from architectures.DenseNet121 import DenseNet121
|
||||||
|
return DenseNet121(device, size)
|
||||||
|
case _:
|
||||||
|
raise ValueError(f"Unknown model: {arch}")
|
||||||
|
|
||||||
|
|
||||||
|
# model architectures
|
||||||
|
from enum import Enum, auto
|
||||||
|
|
||||||
|
class Architecture(Enum):
|
||||||
|
RESNET18 = auto()
|
||||||
|
RESNET50 = auto()
|
||||||
|
INCEPTION = auto()
|
||||||
|
DENSENET121 = auto()
|
||||||
22
architectures/ResNet18.py
Normal file
22
architectures/ResNet18.py
Normal file
@@ -0,0 +1,22 @@
|
|||||||
|
|
||||||
|
import torch.nn as nn
|
||||||
|
from torchvision import models
|
||||||
|
|
||||||
|
# Base model
|
||||||
|
from architectures.Model import Model
|
||||||
|
|
||||||
|
class ResNet18(Model):
|
||||||
|
|
||||||
|
def get(self):
|
||||||
|
m = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
|
||||||
|
|
||||||
|
# freez all layers
|
||||||
|
for param in m.parameters():
|
||||||
|
param.requires_grad = False
|
||||||
|
|
||||||
|
# unfreez the last two
|
||||||
|
for param in m.layer3.parameters(): param.requires_grad = True
|
||||||
|
for param in m.layer4.parameters(): param.requires_grad = True
|
||||||
|
|
||||||
|
m.fc = nn.Linear(m.fc.in_features, self.size)
|
||||||
|
return m.to(self.device)
|
||||||
22
architectures/ResNet50.py
Normal file
22
architectures/ResNet50.py
Normal file
@@ -0,0 +1,22 @@
|
|||||||
|
|
||||||
|
import torch.nn as nn
|
||||||
|
from torchvision import models
|
||||||
|
|
||||||
|
# Base model
|
||||||
|
from architectures.Model import Model
|
||||||
|
|
||||||
|
class ResNet50(Model):
|
||||||
|
|
||||||
|
def get(self):
|
||||||
|
m = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
|
||||||
|
|
||||||
|
# freez all layers
|
||||||
|
for param in m.parameters():
|
||||||
|
param.requires_grad = False
|
||||||
|
|
||||||
|
# unfreez the last two
|
||||||
|
for param in m.layer3.parameters(): param.requires_grad = True
|
||||||
|
for param in m.layer4.parameters(): param.requires_grad = True
|
||||||
|
|
||||||
|
m.fc = nn.Linear(m.fc.in_features, self.size)
|
||||||
|
return m.to(self.device)
|
||||||
5
dependencies.txt
Normal file
5
dependencies.txt
Normal file
@@ -0,0 +1,5 @@
|
|||||||
|
torch
|
||||||
|
torchvision
|
||||||
|
gdown
|
||||||
|
numpy
|
||||||
|
scikit-learn
|
||||||
Reference in New Issue
Block a user