Tensorborad 可视化神经网络模型结构

from  torch.utils.tensorboard import SummaryWriter


writer = SummaryWriter('runs/tensorborad_1.mlp')
writer.add_graph(model,X_train,y_train,X_test)

writer.close()