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

[util] standard checkpoint function naming (#1377)

pull/1382/head
Frank Lee 2022-07-28 09:29:30 +08:00 committed by GitHub
parent 52bc2dc271
commit 0c1a16ea5b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
colossalai/utils/checkpoint/module_checkpoint.py
View File

@ -6,7 +6,7 @@ from colossalai.utils.checkpoint.utils import gather_tensor, scatter_tensor
from typing import Optional
def save_checkpoint(dire: str,
def save_checkpoint(path: str,
epoch: int,
model: torch.nn.Module,
optimizer: Optional[ColossalaiOptimizer] = None,
@ -16,7 +16,7 @@ def save_checkpoint(dire: str,
"""save_checkpoint
save a model, whose parameters are `ColoTensor`s.
Args:
dire (str): directory to save the checkpoint files.
path (str): directory to save the checkpoint files.
epoch (int): the number of epoch
model (torch.nn.Module): a torch module initialized by ColoInitContext
optimizer (ColossalaiOptimizer, optional): optimizers. Defaults to None.
@ -39,7 +39,7 @@ def save_checkpoint(dire: str,
delattr(v, 'save_ready')
# model saving
save_state = {'epoch': epoch, 'model': model_state}
torch.save(save_state, dire + '/epoch_{}_model.pth'.format(epoch), *args, **kwargs)
torch.save(save_state, path + '/epoch_{}_model.pth'.format(epoch), *args, **kwargs)
# delete old dicts
del model_state
@ -57,7 +57,7 @@ def save_checkpoint(dire: str,
if rank == 0:
save_state = {'epoch': epoch, 'optim': optim_state}
torch.save(save_state, dire + '/epoch_{}_optim.pth'.format(epoch), *args, **kwargs)
torch.save(save_state, path + '/epoch_{}_optim.pth'.format(epoch), *args, **kwargs)
# recover colo tensors in rank0
for k, v in optimizer.state_dict()['state'].items():
for n, t in v.items():
@ -71,7 +71,7 @@ def save_checkpoint(dire: str,
dist.barrier()
def load_checkpoint(dire,
def load_checkpoint(path,
epoch: int,
model: torch.nn.Module,
optimizer: Optional[ColossalaiOptimizer] = None,
@ -81,7 +81,7 @@ def load_checkpoint(dire,
"""load_checkpoint
load a model, whose parameters are `ColoTensor`s.
Args:
dire (_type_): _description_
path (_type_): _description_
epoch (int): _description_
rank (int): _description_
model (torch.nn.Module): _description_
@ -96,7 +96,7 @@ def load_checkpoint(dire,
gather_tensor(p)
if rank == 0:
load_state = torch.load(dire + '/epoch_{}_model.pth'.format(epoch), *args, **kwargs)
load_state = torch.load(path + '/epoch_{}_model.pth'.format(epoch), *args, **kwargs)
model.load_state_dict(load_state['model'])
dist.barrier()
@ -118,7 +118,7 @@ def load_checkpoint(dire,
gather_tensor(t)
if rank == 0:
colo_checkpoint = torch.load(dire + '/epoch_{}_optim.pth'.format(epoch), *args, **kwargs)
colo_checkpoint = torch.load(path + '/epoch_{}_optim.pth'.format(epoch), *args, **kwargs)
optimizer.load_state_dict(colo_checkpoint['optim'])
dist.barrier()