具有 max_batch_size > 1 的多 GPU,使用 DDP

作者: xqun3创建于 2023年2月2日更新于 2024年1月9日

def main( config, unet, epoches, text, valid ):

check config path

config_path = Path(config) full_config_path = str(config_path.resolve()) assert config_path.exists(), f'config not found at {full_config_path}' with open(config_path, 'r') as f: config_data = json.loads(f.read()) assert 'checkpoint_path' in config_data, 'checkpoint path not found in config' model_path = Path(config_data['checkpoint_path']) full_model_path = str(model_path.resolve())

setup

内容来源: lucidrains/imagen-pytorch