Descripción del problema
Al escribir un modelo de red neuronal, a menudo es imposible juzgar intuitivamente si el modelo que diseñamos es lo que imaginamos, por lo que es necesario visualizar el modelo.
solución:
1. Instale la capa oculta:
pip install hiddenlayer
2. Ejemplos de aplicación
Por ejemplo, hay una red neuronal de dos capas
class CNN2(nn.Module): #3 layers
def __init__(self,num_classes=10):
super(CNN2, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=3), # 16, 26 ,26
nn.BatchNorm2d(16),
nn.ReLU(inplace=True)
)
self.layer2 = nn.Sequential(
nn.Conv2d(16, 16, kernel_size=1), # 16, 26 ,26
nn.BatchNorm2d(16),
nn.ReLU(inplace=True)
)
self.fc1 = nn.Linear(1*10816, 2048)
self.fc2 = nn.Linear(2048, num_classes)
def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
Funciones personalizadas para imprimir parámetros del modelo y visualizar imágenes
def modelhiddenlayer(model,input,input2):
# import hiddenlayer as h
# params: model = MSDNet(args).cuda() input = (3, 32, 32) input2 = [1 ,3, 32,32]
print(model)
summary(model, input)
vis_graph = h.build_graph(model, torch.zeros(input2).cuda()) # 获取绘制图像的对象
vis_graph.theme = h.graph.THEMES["blue"].copy() # 指定主题颜色
vis_graph.save("./demo1.png") # 保存图像的路径
definir la función principal
def main():
model = CNN2(10).cuda()
input = (1, 28, 28)
input2 = [1,1,28,28]
modelhiddenlayer(model,input,input2)
if __name__ == '__main__':
main()
El resultado se muestra en la figura.
¡éxito!
código completo
import torch
from torch import nn
import torch.nn.functional as F
from torchsummary import summary
import hiddenlayer as h
from torchviz import make_dot
class CNN2(nn.Module): #3 layers
def __init__(self,num_classes=10):
super(CNN2, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=3), # 16, 26 ,26
nn.BatchNorm2d(16),
nn.ReLU(inplace=True)
)
self.layer2 = nn.Sequential(
nn.Conv2d(16, 16, kernel_size=1), # 16, 26 ,26
nn.BatchNorm2d(16),
nn.ReLU(inplace=True)
)
self.fc1 = nn.Linear(1*10816, 2048)
self.fc2 = nn.Linear(2048, num_classes)
def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
def modelhiddenlayer(model,input,input2):
# import hiddenlayer as h
# params: model = MSDNet(args).cuda() input = (3, 32, 32) input2 = [1 ,3, 32,32]
print(model)
summary(model, input)
vis_graph = h.build_graph(model, torch.zeros(input2).cuda()) # 获取绘制图像的对象
vis_graph.theme = h.graph.THEMES["blue"].copy() # 指定主题颜色
vis_graph.save("./demo1.png") # 保存图像的路径
def main():
model = CNN2(10).cuda()
input = (1, 28,28)
input2 = [1 ,1,28,28]
modelhiddenlayer(model,input,input2)
if __name__ == '__main__':
main()