when I build a rnn model ,can batchsize be the forward function parameter?
class BaseModel(nn.Module):
def __init__(self, inputDim, hiddenNum, outputDim, layerNum, cell):
super(BaseModel, self).__init__()
self.hiddenNum = hiddenNum
self.inputDim = inputDim
self.outputDim = outputDim
self.layerNum = layerNum
if cell == "RNN":
self.cell = nn.RNN(input_size=self.inputDim, hidden_size=self.hiddenNum,
num_layers=self.layerNum, dropout=0.0,
nonlinearity="relu", batch_first=True,)
if cell == "LSTM":
self.cell = nn.LSTM(input_size=self.inputDim, hidden_size=self.hiddenNum,
num_layers=self.layerNum, dropout=0.0,
batch_first=True, )
if cell == "GRU":
self.cell = nn.GRU(input_size=self.inputDim, hidden_size=self.hiddenNum,
num_layers=self.layerNum, dropout=0.0,
batch_first=True, )
#print(self.cell)
self.fc = nn.Linear(self.hiddenNum, self.outputDim)
class RNNModel(BaseModel):
def __init__(self, inputDim, hiddenNum, outputDim, layerNum, cell):
super(RNNModel, self).__init__(inputDim, hiddenNum, outputDim, layerNum, cell)
def forward(self, x, batchSize):
h0 = Variable(torch.zeros(self.layerNum*1, batchSize, self.hiddenNum))
rnnOutput, hn = self.cell(x, h0) # rnnOutput 12,20,50 hn 1,20,50
hn = hn.view(batchSize, self.hiddenNum)
fcOutput = self.fc(hn)
return fcOutput