Check "bert-base-uncased" Model Structure
作者:XD / 发表: 2023年4月25日 01:46 / 更新: 2023年4月25日 01:46 / 编程笔记 / 阅读量:1057
Check "bert-base-uncased" Model Structure
import torch
from transformers import BertModel
from torchinfo import summary
model = BertModel.from_pretrained("bert-base-uncased")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
input_size = (1, 512) # input size
input_tensor = torch.randint(0, 10000, input_size, dtype=torch.long, device=device)
summary(model, input_data=input_tensor)
相关标签