Error when using the trained "model.pth" file in transfer learning

36 Views Asked by At

I'm working on a deep learning project. I do not have any problems with training and classification, but when I try to add a section for transfer learning to my project, I get an error. The model I used in training is as follows;

class SimpleCNN(nn.Module):
def __init__(self, num_classes=numClasses):
    super(SimpleCNN, self).__init__()
    self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
    self.relu1 = nn.ReLU()
    self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
    self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
    self.relu2 = nn.ReLU()
    self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
    self.fc1 = nn.Linear(128 * 56 * 56, 512)
    self.relu3 = nn.ReLU()  # Değişiklik burada
    self.dropout = nn.Dropout(0.5)
    self.fc2 = nn.Linear(512, num_classes)

def forward(self, x):
    x = self.conv1(x)
    x = self.relu1(x)
    x = self.pool1(x)
    x = self.conv2(x)
    x = self.relu2(x)
    x = self.pool2(x)
    x = x.view(-1, 128 * 56 * 56)
    x = self.fc1(x)
    x = self.relu3(x)
    x = self.dropout(x)
    x = self.fc2(x)
    return x

I know that the layers of the model I trained and the model I used in transfer learning should be compatible, but I could not achieve this. I will write below the python file I created for transfer learning. Can you help me find out where the problem is?

import json
import torch
import torch.nn as nn
import torchvision.models as models
from torch import optim
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import ImageFolder
from torchvision.transforms import transforms
from model import SimpleCNN

with open("config.json", "r") as config_file:
    config = json.load(config_file)

data_path_main = config["data_path_main"]
input_channels = config["inputChannels"]

alexnet = models.alexnet(pretrained=True)
custom_model = SimpleCNN(num_classes=input_channels)
custom_model.conv1 = alexnet.features[0]

custom_model.fc = nn.Sequential(
    nn.Linear(alexnet.classifier[1].in_features, input_channels)
)

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

dataset = ImageFolder(root=data_path_main, transform=transform)

total_size = len(dataset)
train_size = int(0.7 * total_size)
val_size = int(0.15 * total_size)
test_size = total_size - train_size - val_size

train_dataset, val_dataset, test_dataset = random_split(dataset, [train_size, val_size,     test_size])

batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

model = custom_model
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

num_epochs = 10
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0

    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_loss = running_loss / len(train_loader)
    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss}')

torch.save(model.state_dict(), 't-model.pth')
0

There are 0 best solutions below