Невозможно загрузить модель на устройство Pytorch XLA - PullRequest
Новая
       52

Невозможно загрузить модель на устройство Pytorch XLA

0 голосов
/ 08 апреля 2020

я использовал torch xla и ii попытался отладить мою программу вручную, как показано ниже: 

def fit(model, train_dataset, val_dataset, epochs=1, batch_size=32, warmup_prop=0, lr=5e-5):
    xm.master_print('1')
    device = xm.xla_device()
    xm.master_print('2')
    model = model.to(device)
    xm.master_print('loading train loader')
    train_sampler = DistributedSampler(
        train_dataset,
        num_replicas=xm.xrt_world_size(),
        rank=xm.get_ordinal(),
        shuffle=True
    )



def fit_multiprocessing(rank, flags):
    fit(model, train_dataset, val_dataset, epochs=epochs, batch_size=batch_size, warmup_prop=warmup_prop, lr=lr)

FLAGS = {}
xmp.spawn(fit_multiprocessing, args=(FLAGS,), nprocs=8, start_method='fork')

после печати 1 и 2 процесс обучения (xmp.spawn) застрял навсегда, что означает, что программа не в состоянии завершить sh выполнение этой строки кода: model = model.to (устройство)

вот коды, используемые для создания модели для обучения ТПУ: 

MODEL_PATHS = {
    'xlmr': '../input/xlm-roberta-base/',
}

TRANSFORMERS = {

    'xlmr': (XLMRobertaModel, XLMRobertaTokenizer, XLMRobertaConfig)
}

import torch
from transformers import XLMRobertaConfig,XLMRobertaModel

xlmr_path="/kaggle/input/xlm-roberta-base/"
xlmr = XLMRobertaModel.from_pretrained(xlmr_path)

class Transformer(nn.Module):
    def __init__(self, model, num_classes=1):
        """
        Constructor

        Arguments:
            model {string} -- Transformer to build the model on. Expects "camembert-base".
            num_classes {int} -- Number of classes (default: {1})
        """
        super().__init__()
        self.name = model

        model_class, tokenizer_class, pretrained_weights = TRANSFORMERS[model]

        bert_config = XLMRobertaConfig.from_json_file(MODEL_PATHS[model] + 'config.json')
        bert_config.output_hidden_states = True

        self.transformer = XLMRobertaModel(bert_config)

        self.nb_features = self.transformer.pooler.dense.out_features

        self.pooler = nn.Sequential(
            nn.Linear(self.nb_features, self.nb_features), 
            nn.Tanh(),
        )

        self.logit = nn.Linear(self.nb_features, num_classes)



    def forward(self, tokens):
        """
        Usual torch forward function

        Arguments:
            tokens {torch tensor} -- Sentence tokens

        Returns:
            torch tensor -- Class logits
        """
        _, _, hidden_states = self.transformer(
            tokens, attention_mask=(tokens > 0).long()
        )

        hidden_states = hidden_states[-1][:, 0] # Use the representation of the first token of the last layer

        ft = self.pooler(hidden_states)

        return self.logit(ft)

model = Transformer('xlmr')

извините за не меняя имена переменных, первоначальная настройка была для необработанной модели Bert Base, я просто изменил пути и веса, чтобы вместо этого использовать xlm roberta, я думаю, что в классе трансформатора я что-то упустил, пожалуйста, помогите мне решить эту проблему, спасибо

...