玩具神经网络发生器。
tnng的Python项目详细描述
玩具神经网络发生器
安装
$ pip install tnng
简单模型生成器
^{pr2}$多模模型发生器
#!/usr/bin/env pythonimporttorchimporttorch.nnasnnimporttorchex.nnasexnnfromtnngimportGenerator,MultiHeadLinkedListLayerm=MultiHeadLinkedListLayer()m1=MultiHeadLinkedListLayer()# all layers can be lazy evaluation.m.append([exnn.Linear(64),exnn.Linear(128),exnn.Linear(256)])m.append([nn.ReLU(),nn.ELU()])m.append([exnn.Linear(16),exnn.Linear(32),exnn.Linear(64),])m.append([nn.ReLU(),nn.ELU()])m1.append([exnn.Conv2d(16,1),exnn.Conv2d(32,1),exnn.Conv2d(64,1)])m1.append([nn.MaxPool2d(2),nn.AvgPool2d(2)])m1.append([nn.ReLU(),nn.ELU(),nn.Identity()])m1.append([exnn.Conv2d(32,1),exnn.Conv2d(64,1),exnn.Conv2d(128,1)])m1.append([nn.MaxPool2d(2),nn.AvgPool2d(2)])m1.append([exnn.Flatten(),])m=m+m1m.append([exnn.Linear(128)])m.append([nn.ReLU(),nn.ELU(),nn.Identity()])m.append([exnn.Linear(10)])g=Generator(m)classModel(nn.Module):def__init__(self,idx=0):super(Model,self).__init__()self.model=g[idx]forlayersinself.model:forlayerinlayers:self.add_module(f'{layer}',layer)defforward(self,x,img):forminself.model:iflen(m)==2:ifm[0]isnotNone:x=m[0](x)img=m[1](img)eliflen(m)==1andm[0]isNone:x=torch.cat((x,img),1)else:x=m[0](x)returnxx=torch.randn(128,256)img=torch.randn(128,3,28,28)m=Model()o=m(x,img)print(o.shape)
- 项目
标签: