ropas / pytea Goto Github PK

import torch
import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.cnn = nn.Conv2d(3, 1, 3, 1, 1)
        self.rnn = nn.LSTM(32, 64, 1, batch_first=True)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc = nn.Linear(64, 16)

    def forward(self, x):
        x = self.pool(F.relu(self.cnn(x)))
        x = x.view(-1, 32, 32)
        x, _ = self.rnn(x)
        x = x[:, -1, :].squeeze(1)
        x = F.relu(self.fc(x))
        x = F.softmax(x, dim=-1)
        return x


if __name__ == "__main__":
    net = Model()
    x = torch.randn(2, 3, 64, 64)
    y = net(x)
    target = torch.argmax(torch.randn(2, 16), dim=-1)
    loss = F.cross_entropy(y, target.long())
    loss.backward()
    print(y.size())

   import torch
   import torch.nn as nn

  class Net(nn.Module):
      def __init__(self):
          super(Net, self).__init__()
          self.layers = nn.Sequential(
              nn.Linear(28 * 28, 120),
              nn.ReLU(),
              nn.Linear(80, 10))
  
      def a(self):
          pass

    if __name__ == "__main__":
        n = Net()