| 14 | class DARTSCell(nn.Module): |
| 15 | |
| 16 | def __init__(self, ninp, nhid, dropouth, dropoutx, genotype): |
| 17 | super(DARTSCell, self).__init__() |
| 18 | self.nhid = nhid |
| 19 | self.dropouth = dropouth |
| 20 | self.dropoutx = dropoutx |
| 21 | self.genotype = genotype |
| 22 | |
| 23 | # genotype is None when doing arch search |
| 24 | steps = len(self.genotype.recurrent) if self.genotype is not None else STEPS |
| 25 | self._W0 = nn.Parameter(torch.Tensor(ninp+nhid, 2*nhid).uniform_(-INITRANGE, INITRANGE)) |
| 26 | self._Ws = nn.ParameterList([ |
| 27 | nn.Parameter(torch.Tensor(nhid, 2*nhid).uniform_(-INITRANGE, INITRANGE)) for i in range(steps) |
| 28 | ]) |
| 29 | |
| 30 | def forward(self, inputs, hidden): |
| 31 | T, B = inputs.size(0), inputs.size(1) |