Clean and fix docs

This commit is contained in:
litagin02
2024-03-11 17:05:42 +09:00
parent dada0af2b6
commit 44851a6219
6 changed files with 13 additions and 97 deletions

View File

@@ -215,13 +215,13 @@ def get_logger(
def get_steps(model_path: Union[str, Path]) -> Optional[int]:
"""
モデルのパスからイテレーション番号を取得する
モデルのパスからイテレーション回数を取得する
Args:
model_path (Union[str, Path]): モデルのパス
Returns:
Optional[int]: イテレーション番号
Optional[int]: イテレーション回数
"""
matches = re.findall(r"\d+", model_path) # type: ignore

View File

@@ -27,7 +27,7 @@ def load_checkpoint(
for_infer (bool): 推論用に読み込むかどうかのフラグ
Returns:
tuple[torch.nn.Module, Optional[torch.optim.Optimizer], float, int]: 更新されたモデルとオプティマイザー、学習率、イテレーション番号
tuple[torch.nn.Module, Optional[torch.optim.Optimizer], float, int]: 更新されたモデルとオプティマイザー、学習率、イテレーション回数
"""
assert os.path.isfile(checkpoint_path)
@@ -104,7 +104,7 @@ def save_checkpoint(
model (torch.nn.Module): 保存するモデル
optimizer (Union[torch.optim.Optimizer, torch.optim.AdamW]): 保存するオプティマイザー
learning_rate (float): 学習率
iteration (int): イテレーション数
iteration (int): イテレーション
checkpoint_path (Union[str, Path]): 保存先のパス
"""
logger.info(

View File

@@ -22,7 +22,7 @@ def load_safetensors(
for_infer (bool): 推論用に読み込むかどうかのフラグ
Returns:
tuple[torch.nn.Module, Optional[int]]: 読み込まれたモデルとイテレーション番号(存在する場合)
tuple[torch.nn.Module, Optional[int]]: 読み込まれたモデルとイテレーション回数(存在する場合)
"""
tensors: dict[str, Any] = {}
@@ -64,7 +64,7 @@ def save_safetensors(
Args:
model (torch.nn.Module): 保存するモデル
iteration (int): イテレーション番号
iteration (int): イテレーション回数
checkpoint_path (Union[str, Path]): 保存先のパス
is_half (bool): モデルを半精度で保存するかどうかのフラグ
for_infer (bool): 推論用に保存するかどうかのフラグ