colossalai
26 строк · 833.0 Байт
1import torch
2import torch.nn as nn
3
4from colossalai.shardformer.shard.utils import set_tensors_to_none
5
6
7class Net(nn.Module):
8def __init__(self) -> None:
9super().__init__()
10self.layers = nn.Sequential(nn.Linear(1, 2), nn.Linear(2, 3))
11self.out = nn.Linear(3, 1)
12
13
14def test_release_layer():
15orig_cuda_allocated = torch.cuda.memory_allocated()
16model = Net().cuda()
17set_tensors_to_none(model, exclude={model.layers[0]})
18assert model.layers[1].weight is None
19assert model.layers[1].bias is None
20assert model.out.weight is None
21assert model.out.bias is None
22set_tensors_to_none(model)
23assert model.layers[0].weight is None
24assert model.layers[0].bias is None
25assert len(list(model.parameters())) == 0
26assert torch.cuda.memory_allocated() == orig_cuda_allocated
27