Обучение экстрактора пользовательских функций в стабильных базовых условиях3. Начиная с предварительно обученных весов? ⇐ Python

[Anonymous]

Я использую следующий экстрактор пользовательских функций для своей модели StableBaselines3:

Код: Выделить всё

import torch.nn as nn
from stable_baselines3 import PPO

class Encoder(nn.Module):
def __init__(self, input_dim, embedding_dim, hidden_dim, output_dim=2):
super(Encoder, self).__init__()
self.encoder = nn.Sequential(
nn.Linear(input_dim, embedding_dim),
nn.ReLU()
)
self.regressor = nn.Sequential(
nn.Linear(embedding_dim, hidden_dim),
nn.ReLU(),
)

def forward(self, x):
x = self.encoder(x)
x = self.regressor(x)
return x

model = Encoder(input_dim, embedding_dim, hidden_dim)
model.load_state_dict(torch.load('trained_model.pth'))

class CustomFeatureExtractor(BaseFeaturesExtractor):
def __init__(self, observation_space, features_dim):
super(CustomFeatureExtractor, self).__init__(observation_space, features_dim)
self.model = model  # Use the pre-trained model as the feature extractor

self._features_dim = features_dim

def forward(self, observations):
features = self.model(observations)
return features

policy_kwargs = {
"features_extractor_class": CustomFeatureExtractor,
"features_extractor_kwargs": {"features_dim": 64}
}

model = PPO("MlpPolicy", env=envs, policy_kwargs=policy_kwargs)

# Freeze all layers
for param in model.parameters():
param.requires_grad = False

На данный момент модель хорошо обучена, без проблем и с хорошими результатами. Теперь я хочу не замораживать веса и попытаться обучить Feature Extractor, начиная с начального предварительно обученного веса. Как я могу сделать это с помощью такого специального экстрактора функций, определенного как класс внутри другого класса? Мой экстрактор функций отличается от описанного в документации, поэтому я не уверен, что он будет обучен. Или он начнет обучение, если я разморозлю слои?

Подробнее здесь: https://stackoverflow.com/questions/787 ... -pre-train