layout dm

LayoutDM: modelo de difusão discreta para geração de layout controlável (CVPR2023)

Este repositório é uma implementação oficial do artigo intitulado acima. Consulte a página do projeto ou documento para obter mais detalhes.

Configurar

Aqui descrevemos a configuração necessária para o treinamento e avaliação do modelo.

Requisitos

Verificamos a reprodutibilidade neste ambiente.

• Python3.7
• CUDA11.3
• PyTorch 1.10

Recomendamos usar Poetry (todas as configurações e dependências em pyproject.toml). A geometria Pytorch fornece roda de pré-construção independente para uma combinação da versão PyTorch e CUDA (consulte PyG:Instalação para obter detalhes). Se o seu ambiente não corresponder ao acima, atualize as dependências.

Como instalar

1. Instale poesia (veja documentos oficiais). Recomendamos fazer um virtualenv e instalar poesia dentro dele.

curl -sSL https://install.python-poetry.org | python3 -

2. Instale dependências (pode ser lento ..)

poetry install

3. Baixe recursos e descompacte

wget https://github.com/CyberAgentAILab/layout-dm/releases/download/v1.0.0/layoutdm_starter.zip
unzip layoutdm_starter.zip

Os dados são descompactados na seguinte estrutura:

 download
- clustering_weights
- datasets
- fid_weights
- pretrained_weights

Experimentar

Importante : encontramos alguns erros críticos que não podem ser corrigidos rapidamente ao usar múltiplas GPUs. Defina CUDA_VISIBLE_DEVICES=<GPU_ID> para forçar o modelo a usar uma única GPU.

Nota: nossa estrutura principal é baseada na hidra. É conveniente lidar com dezenas de argumentos hierarquicamente, mas pode exigir alguns esforços adicionais se alguém for novo no Hydra.

Demonstração

Execute um notebook jupyter em notebooks/demo.ipynb. Você pode obter e renderizar os resultados de seis tarefas de geração de layout em dois conjuntos de dados (Rico e PubLayNet).

Treinamento

Você também pode treinar seu próprio modelo do zero, por exemplo,

bash bin/train.sh rico25 layoutdm

, onde o primeiro e o segundo argumentos especificam o conjunto de dados (escolhas) e o tipo de experimento (escolhas), respectivamente. Observe que para treinamento/teste, o estilo dos argumentos é key=value porque usamos hydra, ao contrário do popular --key value (por exemplo, argparse).

Teste

poetry run python3 -m src.trainer.trainer.test 
    cond= < COND > 
    job_dir= < JOB_DIR > 
    result_dir= < RESULT_DIR > 
    < ADDITIONAL_ARGS >

<COND> pode ser: (incondicional, c, cwh, parcial, refinamento, relação)

Por exemplo, se você deseja testar o modelo LayoutDM fornecido em C->S+P , o comando é o seguinte:

 poetry run python3 -m src.trainer.trainer.test cond=c dataset_dir=./download/datasets job_dir=./download/pretrained_weights/layoutdm_rico result_dir=tmp/dummy_results

Consulte TestConfig para mais opções disponíveis. Abaixo estão algumas opções populares para <ADDITIONAL_ARGS>

• is_validation=true : usado para avaliar o desempenho da geração no conjunto de validação em vez do conjunto de teste. Isso deve ser usado ao ajustar os hiperparâmetros.
• sampling=top_p top_p=<TOP_P> : use amostragem top-p com p=<TOP_P> em vez da amostragem padrão.

Avaliação

poetry run python3 eval.py < RESULT_DIR > 

Citação

Se você achar este código útil para sua pesquisa, cite nosso artigo:

 @inproceedings{inoue2023layout,
  title={{LayoutDM: Discrete Diffusion Model for Controllable Layout Generation}},
  author={Naoto Inoue and Kotaro Kikuchi and Edgar Simo-Serra and Mayu Otani and Kota Yamaguchi},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2023},
  pages={10167-10176},
}