mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-07-12 19:06:17 +00:00
2053e138a2
* [CLI] add CLI launcher
* Revert "[CLI] add CLI launcher"
This reverts commit df7e6506d4.
* [context]use meta tensor to init model lazily.
* polish
* make module with device kwargs bypass the normal init.
* change unit test to adapt updated context.
23 lines
585 B
Python
23 lines
585 B
Python
import torch
|
|
from colossalai.utils.model.lazy_init_context import LazyInitContext
|
|
from torchvision.models import resnet34
|
|
|
|
|
|
def test_lazy_init():
|
|
ctx = LazyInitContext()
|
|
with ctx:
|
|
model = resnet34(num_classes=10)
|
|
for param in model.parameters():
|
|
assert param.is_meta
|
|
for buffer in model.buffers():
|
|
assert buffer.is_meta
|
|
ctx.lazy_init_parameters(model)
|
|
for param in model.parameters():
|
|
assert not param.is_meta
|
|
for buffer in model.buffers():
|
|
assert not buffer.is_meta
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_lazy_init()
|