net.py 2.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. import torch
  2. import torch.nn as nn
  3. #定义神经网络
  4. class TextClassifier(nn.Module):
  5. def __init__(self, vocab_size,embed_dim,num_classes,padding_idx):
  6. super().__init__()
  7. #定义一个词嵌入层
  8. #embedding层的大小是vocab_size*embed_dim,表示词表大小乘以词向量的维度
  9. # 5000*128
  10. self.embedding = nn.Embedding(vocab_size,embed_dim,padding_idx=padding_idx)
  11. #定义一个线性层
  12. #线性层的大小是 embed_dim*num_classes,表示词向量维度乘以类别数量
  13. #例如 128*4
  14. self.fc = nn.Sequential(
  15. nn.Linear(embed_dim,128),
  16. nn.BatchNorm1d(128),
  17. nn.ReLU(),
  18. nn.Linear(128,64),
  19. nn.BatchNorm1d(64),
  20. nn.ReLU(),
  21. nn.Linear(64,32),
  22. nn.BatchNorm1d(32),
  23. nn.ReLU(),
  24. nn.Linear(32,num_classes)
  25. )
  26. #例如,输入一个2*6的文本序列,表示2个样本,每个样本6个单词
  27. def forward(self,x):
  28. #将每个单词都转为词向量
  29. #例如,输入的x 是2*6的,每个单词用128维表示
  30. #那么x经过embedding层后,就会变为2*6*128的结果
  31. # b ,s, e = x.shape[0],x.shape[1],self.embed_dim
  32. x = self.embedding(x)
  33. #计算一个句子中各个单词的词向量和
  34. #也就是将每个句子中的单词的词向量相加后,得到每个句子的向量表示
  35. #例如, 2*6*128的x ,就会被转换为2*128的结果
  36. #也就是2个句子,每个句子都用128维来表示
  37. x = torch.sum(x,dim=1)
  38. #最终的分类结果
  39. x=self.fc(x)
  40. return x
  41. if __name__== "__main__":
  42. #定义一个词表大小为5000,词向量维度为128,类别数量为4,填充索引为0的文本分类器
  43. model = TextClassifier(vocab_size=5000,
  44. embed_dim=128,
  45. num_classes=4,
  46. padding_idx=0)
  47. #打印模型结构
  48. print(model)
  49. #定义一个2*6的输入数据,表示2个样本,每个样本6个单词
  50. x = torch.randint(0,5000,(2,6))
  51. #打印输入数据
  52. print(x)
  53. #将输入数据传入模型中,得到输出结果
  54. output = model(x)
  55. #打印输出结果
  56. print(output)
  57. print(output.shape)