PyG Temporal: создание наборов данных из небольших синтетических данных

Question

Общий

В отличие от примеров, приведенных в геометрических временных документах pytorch , где используются большие временные ряды, я пытаюсь обучить модель на большом количестве статических сетей с небольшим отставанием (обычно <10). До сих пор мне удалось обучить модель, но то, как я это сделал, кажется хакерским, и у меня возникают проблемы с памятью. Я также не смог тренироваться на графическом процессоре.

Мои данные состоят из набора (на данный момент суперпростого правила большинства) динамики, запущенной в случайно сгенерированной сети Эрдёша-Реньи для заданного количества «временных шагов», хранящихся в массиве формы (количество узлов, отставание). Затем я преобразовываю каждый набор данных в сигнал, совместимый с PyG, и передаю его через модель.

Реализация

Моей основной точкой входа в программу являетсяExperimentучебный класс.

Генерация и преобразование данных

Здесь Datasetэто просто dataclassкоторый содержит как данные о характеристиках узла, так и информацию об уровне графика (фактический график), который используется позже при генерации сигналов.

Преобразованием данных занимается SignalTransformкласс, который возвращает список геометрических временных StaticGraphTemporalSignalобъекты (класс здесь ).

      class Experiment:
   [...]
    def generate_data(self) -> None:
        """Populates the `self.dataset` attr with `Dataset` objects
        containing the dynamics ran on each network. Datasets have shape
        (samples, nodes, lagsteps) and also contain network info.
        """
        self.setup()
        for nw in range(self.cfg.exp.num_networks):
            dataset = Dataset(self.networks[nw], self.cfg)
            self.set_dynamics(self.networks[nw])
            for sample in range(self.cfg.exp.num_samples):
                x = self.dynamics.initial_state()
                # import pdb; pdb.set_trace()
                dataset.data[sample, :, 0] = x
                for lagstep in range(1, self.cfg.exp.lag):
                    if lagstep % self.cfg.exp.lagstep == 0:
                        dataset.data[sample, :, lagstep] = x = self.dynamics.step(x)
            self.dataset.append(dataset)

    def generate_signals(self) -> None:
        self.signals = dict()
        for nw_idx in range(self.cfg.exp.num_networks):
            self.signals[nw_idx] = SignalTransform(self.dataset, nw_idx).get_signals()

   [...]

Модель

Я в значительной степени использовал модель и методологию, описанные в документах и ​​этой записной книжке , для обучения. Файл, содержащий всю мою методику тренировок, можно найти здесь .

Короче говоря, я создал Temporal GNN с допустимыми (надеюсь) параметрами для моего случая и запустил фактическое обучение и оценку модели в TestModelучебный класс:

      class TestModel:
    [...]
    def train(self):
        optimizer = torch.optim.Adam(self.model.parameters(), lr=0.01)
        self.model.train()
        for epoch in range(10):
            loss = 0
            step = 0
            for signal in self.train_set:
                for snapshot in signal:
                    snapshot = snapshot.to(self.device)
                    y_hat = self.model(snapshot.x, snapshot.edge_index)
                    loss = loss + torch.mean((y_hat - snapshot.y) ** 2)
                    step += 1
            loss = loss / (step + 1)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            print("Epoch {} train MSE: {:.4f}".format(epoch, loss.item()))
    [...]

Проблемы и конечные цели

Я хотел бы выполнить какую-то пакетную обработку, чтобы я мог тренироваться на больших сетях/наборах данных. Текущая реализация на наборе из 1000 сетей с 500 узлами и 5 шагами задержки в каждой требует более 10 ГБ ОЗУ. Я совершенно не уверен, как действовать (или даже если это вообще возможно, я совершенно новичок в ML tbh). Единственный способ, который я могу придумать, - это фактически «распаковать» снимки как pytorch_geometric.Dataобъекты и создание наборов данных и подобных пакетов. Я также хотел бы иметь возможность тренироваться с помощью cuda.

Моей конечной целью будет классификация на уровне графа (т. е. входная сеть достигает состояния «стабильное — консенсус»), но я думаю, что это не составит труда, если вышеуказанные проблемы будут решены.

Любая помощь будет принята с благодарностью - спасибо :)

PS Извините за количество внешних ссылок, но я постарался, чтобы этот вопрос был как можно короче. Не уверен, что это нормально.