Сохраните обученную модель автоэнкодера и загрузите ее в другой модуль в Pytorch. - PullRequest
Новая
       53

Сохраните обученную модель автоэнкодера и загрузите ее в другой модуль в Pytorch.

0 голосов
/ 01 августа 2020

У меня есть autoencoder, который я обучил, и теперь я хочу сохранить лучшую модель. Я хотел бы загрузить эту обученную модель в другой блокнот Jupiter, а затем использовать ее в режиме вывода. Это мой autoencoder класс: 

class DEC_AE(nn.Module):
"""
DEC auto encoder - this class defines the architecture of our Auto-encoder
"""

def __init__(self, num_classes, num_features): #what are number of classes (10) and number of features (10)?
    super(DEC_AE, self).__init__()
    self.dropout = nn.Dropout(p=0.1)#using dropout 0.1-is it according to the paper?
    self.fc1 = nn.Linear(28 * 28, 500)
    self.fc2 = nn.Linear(500, 500)
    self.fc3 = nn.Linear(500, 2000)
    self.fc4 = nn.Linear(2000, num_features)#the architecture used is exactly like the one in the paper
    self.relu = nn.ReLU()
    self.fc_d1 = nn.Linear(500, 28 * 28)
    self.fc_d2 = nn.Linear(500, 500)
    self.fc_d3 = nn.Linear(2000, 500)
    self.fc_d4 = nn.Linear(num_features, 2000)
    self.alpha = 1.0 #what is alpha?
    
    self.clusterCenter = nn.Parameter(torch.zeros(num_classes, num_features))#nn-Parameter
    self.pretrainMode = True
    #here we initialize all the weights
    for m in self.modules():
        if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
            torch.nn.init.xavier_uniform_(m.weight)

def setPretrain(self, mode):
    """To set training mode to pretrain or not,
    so that it can control to run only the Encoder or Encoder+Decoder"""
    self.pretrainMode = mode

#this one I need
def forward(self, x):#a batch of images
    x = x.view(-1, 1 * 28 * 28)
    x = self.dropout(x)
    x = self.fc1(x)
    x = self.relu(x)
    x = self.fc2(x)
    x = self.relu(x)
    x = self.fc3(x)
    x = self.relu(x)
    x = self.fc4(x)
    x_e = x #the output of encoder

    # if not in pre_train mode, we need encoder and t distribution output
    if self.pretrainMode is False:
        return x, F.softmax(x_e, dim=1)
        #why softmax of x?

    # encoder is done, followed by decoder
    x = self.fc_d4(x)
    x = self.relu(x)
    x = self.fc_d3(x)
    x = self.relu(x)
    x = self.fc_d2(x)
    x = self.relu(x)
    x = self.fc_d1(x)
    x_de = x.view(-1, 1, 28, 28)
    
    return x_e, x_de

, и я сохраняю его вот так 

 torch.save(dec_ae, 'bestModel'.format(best_acc))

, а затем в другом файле я пытаюсь загрузить файл следующим образом: 

def load_checkpoint(filepath):
   checkpoint = torch.load(filepath)
   model = checkpoint['model']
   model.load_state_dict(checkpoint['state_dict'])
   for parameter in model.parameters():
      parameter.requires_grad = False

  model.eval()
  return model

model = load_checkpoint('bestModel')
model.cuda()

Но я получаю эту ошибку: 

 AttributeError: Can't get attribute 'DEC_AE' on <module '__main__'>

Я прочитал в сообщении, что мне нужно импортировать модуль следующим образом: 

from Autoencoder import DEC_AE

, но он не работает. Есть идеи, как это решить? спасибо!

...