mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-05-02 13:45:36 +00:00
250be4d31e
* [utils] integrated colotensor with lazy init context * polish code * polish code * polish code
52 lines
1.3 KiB
Python
52 lines
1.3 KiB
Python
import torch
|
|
from colossalai.utils.model.lazy_init_context import LazyInitContext
|
|
from torchvision.models import resnet34
|
|
import random
|
|
import numpy as np
|
|
|
|
MANUAL_SEED = 0
|
|
random.seed(MANUAL_SEED)
|
|
np.random.seed(MANUAL_SEED)
|
|
torch.manual_seed(MANUAL_SEED)
|
|
|
|
|
|
def test_lazy_init_with_meta():
|
|
ctx = LazyInitContext(to_meta=True)
|
|
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 name, param in model.named_parameters():
|
|
assert not param.is_meta, name
|
|
|
|
for buffer in model.buffers():
|
|
assert not buffer.is_meta
|
|
|
|
|
|
def test_lazy_init_without_meta():
|
|
ctx = LazyInitContext(to_meta=False)
|
|
with ctx:
|
|
model = resnet34(num_classes=10)
|
|
|
|
for param in model.parameters():
|
|
assert not param.is_meta
|
|
for buffer in model.buffers():
|
|
assert not buffer.is_meta
|
|
|
|
conv1_weight_before_init = model.conv1.weight.clone()
|
|
ctx.lazy_init_parameters(model)
|
|
conv1_weight_after_init = model.conv1.weight.clone()
|
|
|
|
assert not torch.allclose(conv1_weight_after_init, conv1_weight_before_init)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_lazy_init_with_meta()
|
|
test_lazy_init_without_meta()
|