-
Notifications
You must be signed in to change notification settings - Fork 73
/
test.py
69 lines (44 loc) · 1.43 KB
/
test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import warnings
from torch.utils.data.dataloader import DataLoader
from dataloaders.datasets.cityscapes import CityscapesSegmentation
from config_utils.search_args import obtain_search_args
from utils.loss import SegmentationLosses
import torch
import numpy as np
from auto_deeplab import AutoDeeplab
model = AutoDeeplab(19, 12).cuda()
args = obtain_search_args()
args.cuda = True
criterion = SegmentationLosses(weight=None, cuda=args.cuda).build_loss(mode=args.loss_type)
def save_grad(name):
def hook(grad):
grads[name] = grad
return hook
args.crop_size = 64
dataset = CityscapesSegmentation(args, r'E:\BaiduNetdiskDownload\cityscapes', 'train')
loader = DataLoader(dataset, batch_size=2, shuffle=True)
grads = {}
def save_grad(name):
def hook(grad):
grads[name] = grad
return hook
for i, sample in enumerate(loader):
image, label = sample['image'].cuda(), sample['label'].cuda()
# from thop import profile
# params, flops = profile(model, inputs=(image, ))
# print(params)
# print(flops)
model.betas.register_hook(save_grad('y'))
prediction = model(image)
# y = 1e-3*torch.randn(12, 4, 3).cuda()
# criterion = torch.nn.MSELoss()
# z = criterion(prediction, label)
z = prediction.mean()
z.backward()
print(grads['y'])
print(grads['y'].shape)
# print(grads['y1'])
# print(grads['y1'].shape)
if i == 0:
exit()
# 查看 y 的梯度值