1、对于自己网络有外部参数需要输入时:如下所示,外部参数有:nstack, inp_dim, oup_dim ,bn
class PoseNet(nn.Module):
def __init__(self, nstack, inp_dim, oup_dim, bn=False, increase=128, init_weights=True, **kwargs):
2、定义可视化结构
1)首先,对网络PoseNet传入参数输入:如,nstack=4; inp_dim=256; oup_dim =54;bn = True
2)其次,定义输入。input =torch.rand(1,128,128,3)
3)最后,在如下代码的comment 处写上pose
with SummaryWriter(comment='pose') as w:
w.add_graph(pose, (input,))
if __name__ == '__main__':
from time import time
pose = PoseNet(4, 256, 54, bn=True) # .cuda()
for param in pose.parameters():
if param.requires_grad:
print('param autograd')
break
t0 = time()
input = torch.rand(1, 128, 128, 3) # .cuda()
print(pose)
output = pose(input) # type: torch.Tensor
output[0][0].sum().backward()
with SummaryWriter(comment='pose') as w:
w.add_graph(pose, (input,))
t1 = time()