Учебник 3: Модель

В этом учебнике мы кратко рассмотрим модели генерации танца, содержащиеся в XRMoGen, и их интерфейс.

Модели генерации танца

На данный момент XRMoGen содержит два алгоритма генерации танца:

• Bailando: Siyao et al., Bailando: 3D Dance Generation by Actor-Critic GPT with Choreographic Memory, CVPR 2022;
• DanceRevolution: Huang et al., Dance Revolution: Long-Term Dance Generation with Music via Curriculum Learning, ICLR 2021.

В обеих моделях реализованы интерфейсы для обучения и проверки/тестирования раннера:

@DANCE_MODELS.register_module()
class MyDanceModel(nn.Module):
    ...
    def train_step(self, data, optimizer, **kwargs):
        ...
    def val_step(self, data, optimizer=None, **kwargs):
        ...

Входные данные data представляют собой временно связанные музыкальные функции и последовательность поз человека в 3D. Для обучения используются как музыкальные функции, так и 3D позы для контролируемого обучения. Для тестирования используется только музыкальная функция с начальной позой (код) для генерации танца.

Результатом проверки является словарь, где output_pose — это сгенерированный танец размером (nframes, njoints=24, ndim=3). file_name — строка имени файла соответствующей выходной позы.

outputs = {
    'output_pose': results[0],
    'file_name': data['file_names'][0]
}

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