github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish colossalai/cli/benchmark/benchmark.py code style (#2287)

pull/2317/head
アマデウス 2023-01-03 18:23:11 +08:00 committed by Frank Lee
parent 1c29b173c9
commit 49715a78f0
1 changed files with 8 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
colossalai/cli/benchmark/benchmark.py
View File

@ -1,16 +1,17 @@
import colossalai
from functools import partial
from typing import Dict, List
import click
import torch.multiprocessing as mp
from functools import partial
from typing import List, Dict
import colossalai
from colossalai.cli.benchmark.utils import find_all_configs, get_batch_data, profile_model
from colossalai.context import Config
from colossalai.context.random import reset_seeds
from colossalai.core import global_context as gpc
from colossalai.logging import disable_existing_loggers, get_dist_logger
from colossalai.utils import free_port, MultiTimer
from colossalai.cli.benchmark.utils import find_all_configs, profile_model, get_batch_data
from colossalai.utils import MultiTimer, free_port
from .models import MLP
@ -53,7 +54,7 @@ def run_dist_profiling(rank: int, world_size: int, port_list: List[int], config_
port_list (List[int]): a list of free ports for initializing distributed networks
config_list (List[Dict]): a list of configuration
hyperparams (Config): the hyperparameters given by the user
"""
# disable logging for clean output