Ошибка при загрузке state_dict для привычной модели

Question

У меня были проблемы при загрузке весов модели. Вот некоторые детали модели

class InceptionV4(nn.Module):

   def __init__(self, num_classes=1001):
       super(InceptionV4, self).__init__()
       # Special attributs
       self.input_space = None
       self.input_size = (299, 299, 3)
       self.mean = None
       self.std = None
       # Modules
       self.features = nn.Sequential(
           BasicConv2d(3, 32, kernel_size=3, stride=2),
           BasicConv2d(32, 32, kernel_size=3, stride=1),
           BasicConv2d(32, 64, kernel_size=3, stride=1, padding=1),
           Mixed_3a(),
           Mixed_4a(),
           Mixed_5a(),
           Inception_A(),
           Inception_A(),
           Inception_A(),
           ...
       )
       self.avg_pool = nn.AvgPool2d(8, count_include_pad=False)
       self.last_linear = nn.Linear(1536, num_classes)

Я попытался сохранить вес, что-то вроде torch.save(model.state_dict(), weight_name), а затем снова перезагрузить model.load_state_dict(torch.load(weight_name)) но получил эти ошибки:

Missing key(s) in state_dict: "features.0.conv.weight", "features.0.bn.weight", "features.0.bn.bias", "features.0.bn.running_mean", "features.0.bn.running_var", "features.1.conv.weight", "features.1.bn.weight", "features.1.bn.bias", "features.1.bn.running_mean", "features.1.bn.running_var", "features.2.conv.weight", "features.2.bn.weight

, а также:

Unexpected key(s) in state_dict: "conv.0.conv1.0.weight", "conv.0.conv1.0.bias", "conv.0.conv1.2.weight", "conv.0.conv1.2.bias", "conv.0.conv1.2.running_mean", "conv.0.conv1.2.running_var", "conv.0.conv1.2.num_batches_tracked", "conv.0.conv2.0.weight", "conv.0.conv2.0.bias", "conv.0.conv2.2.weight", "conv.0.conv2.2.bias", "conv.0.conv2.2.running_mean", "conv.0.conv2.2.running_var", "conv.0.conv2.2.num_batches_tracked", "conv.1.conv1.0.weight", "conv.1.conv1.0.bias", "conv.1.conv1.2.weight", "conv.1.conv1.2.bias", "conv.1.conv1.2.running_mean", "conv.1.conv1.2.running_var", "conv.1.conv1.2.num_batches_tracked

Есть намеки на это? Заранее спасибо.

Answer 1

Я сталкивался с этой проблемой несколько раз. Ошибка указывает, что ваша модель state_dict имеет имена, отличные от загруженной pre-trained weights.

Я не вижу предварительно обученную модель для Inception_v4 в torchvision модельном зоопарке, поэтому было бы немного трудно точно определить, где у вашего класса InceptionV4 есть проблема с несовпадающим диктом.

Независимо от того, где вы получили файл pre-trained, но ключевой момент - определить вашу модель так же, как код модели pre-trained, и вы можете плавно загрузить файл весов.

Вот некоторые показатели, по которым ваш код отличается от модели:

# change self.features -> self.conv: This helps in solving mismatched names.

self.conv = nn.Sequential(...)


# Google how to change the BatchNorm in your current pytorch version 
# and  the older pytorch version which the pretrained model was defined.

conv.1.conv1.2.num_batches_tracked  # it is deprecated in pytorch version 0.4 or newer

Подсказка:

# Define your model (or parts you want to reuse) the same as the original 

Надеюсь, это поможет:)