mirror of https://github.com/hpcaitech/ColossalAI
[hotfix] rm test_tensor_detector.py (#413)
parent
370f567e7d
commit
a37bf1bc42
|
@ -1,41 +0,0 @@
|
||||||
#!/usr/bin/env python
|
|
||||||
# -*- encoding: utf-8 -*-
|
|
||||||
|
|
||||||
import torch
|
|
||||||
import torch.nn as nn
|
|
||||||
|
|
||||||
from colossalai.utils import TensorDetector
|
|
||||||
|
|
||||||
class MLP(nn.Module):
|
|
||||||
def __init__(self):
|
|
||||||
super().__init__()
|
|
||||||
self.mlp = nn.Sequential(nn.Linear(64, 8),
|
|
||||||
nn.ReLU(),
|
|
||||||
nn.Linear(8, 32))
|
|
||||||
|
|
||||||
def forward(self, x):
|
|
||||||
return self.mlp(x)
|
|
||||||
|
|
||||||
def test_tensor_detect():
|
|
||||||
|
|
||||||
data = torch.rand(64, requires_grad=True).cuda()
|
|
||||||
data.retain_grad()
|
|
||||||
model = MLP().cuda()
|
|
||||||
|
|
||||||
detector = TensorDetector(log='test', include_cpu=False, module=model)
|
|
||||||
|
|
||||||
detector.detect()
|
|
||||||
out = model(data)
|
|
||||||
|
|
||||||
detector.detect()
|
|
||||||
loss = out.sum()
|
|
||||||
detector.detect()
|
|
||||||
loss.backward()
|
|
||||||
detector.detect()
|
|
||||||
model = MLP().cuda()
|
|
||||||
detector.detect()
|
|
||||||
detector.close()
|
|
||||||
torch.cuda.empty_cache()
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
test_tensor_detect()
|
|
Loading…
Reference in New Issue