画pytorch模型图,以及参数计算的方法

刚入pytorch的坑，代码还没看太懂。之前用keras用习惯了，第一次使用pytorch还有些不适应，希望广大老司机多多指教。首先说说，我们如何可视化模型。在keras中就一句话，keras.summary()，或者plot_model()，就可以把模型展现的淋漓尽致。但是pytorch中好像没有这样一个api让我们直观的看到模型的样子。但是有网友提供了一段代码，可以把模型画出来，对我来说简直就是如有神助啊。就是conv -> relu -> maxpool -> conv -> relu -> maxpool -> fc大家在可视化的时候，直接复制make_dot那段代码即可，然后需要初始化一个net，以及这个网络需要的数据规模，此处就以 这段代码为例，初始化一个模型net，准备这个模型的输入数据x，shape为 然后把数据传入模型得到输出结果y。传入make_dot即可得到下图。最后输出该网络的各种参数。