#模型研发参数查看#
在进行模型研发或者模型设计时需要对模型结构和模型模型参数进行统计,pytorch提供了相应的统计工具summary,如下为使用代码:
import torch
import torch.nn as nn
from torchsummary import summary
class VisionTransformer(nn.Module):
...
model = VisionTransformer(
**kwargs,
)
summary(model,input_size=(MODEL_INPUT_SIZE))
此时通过屏幕即可打印相关模型信息,包含有每一层的信息、可训练参数、总参数等:
但是会有一些报错信息用于记录:
错误一
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same
原因分析:模型和输入的tensor不在同一个device上,一个在CPU一个在GPU上。
def summary(model, input_size, batch_size=-1, device="cuda"):
...
if device == "cuda" and torch.cuda.is_available():
dtype = torch.cuda.FloatTensor
else:
dtype = torch.FloatTensor
...
x = [torch.rand(2, *in_size).type(dtype) for in_size in input_size]
进入 summary函数后可以看到有一个参数'device' ,这个参数会将data传入到相应的设备上,而此时model却是在CPU上。所以,weights在cpu上,data在GPU上。有两种解决方案,即将model放在GPU或者data放在CPU。(推荐都放在GPU上)
#方法一,将data放在cpu
summary(model.cuda(),input_size=( 3, 384, 384),device="cpu")
#方法二,将模型放在GPU
# vit某些算子CPU不支持的,所以推荐方法二
summary(model.cuda(),input_size=( 3, 384, 384))
错误二
维度不正确,这类错误就比较多了,比如你的batch_size没有对应上,或者(c,w,h)维度没有和你的模型输入正确对应上,都会产生错误,需要仔细和模型输入进行对比即可。