특정 Layer의 weight를 고정시키기

import torch
import torch.nn as nn
import torch.optim as optim

# Load a pre-trained model
model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True)

# Fix the weights of the first 5 layers
for param in model.parameters():
    param.requires_grad = False

for param in model.layer4.parameters():
    param.requires_grad = True

# Define a new classifier
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 2)

# Define a loss function
criterion = nn.CrossEntropyLoss()

# Define an optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Train the model
for epoch in range(num_epochs):
    # Forward pass
    outputs = model(inputs)
    loss = criterion(outputs, labels)

    # Backward pass for unfrozen layers
    optimizer.zero_grad()
    loss.backward()

    # Update the weights for unfrozen layers
    optimizer.step()