How to build basic Neural Network with PyTorch

The nn package in PyTorch provides high level abstraction for building neural networks. In this post we will build a simple Neural Network using PyTorch nn package.

Import torch and define layers dimensions



import torch

batch_size, input_dim, hidden_dim, out_dim = 32, 100, 100, 10

Create input, output tensors



input_tensor = torch.randn(batch_size, input_dim)
output_tensor = torch.randn(batch_size, out_dim)

Define model using nn package



model = torch.nn.Sequential(
    torch.nn.Linear(input_dim, hidden_dim),
    torch.nn.Tanh(),
    torch.nn.Linear(hidden_dim, out_dim),
)

Define loss function, optimizer and learning rate



loss_function = torch.nn.MSELoss(reduction='sum')

lr = 1e-5

sgd_optimizer = torch.optim.SGD(model.parameters(), lr=lr)


Train the model



for i in range(200):
    predicted_value = model(input_tensor)
    loss = loss_function(predicted_value, output_tensor)
    print(i, loss.item())

    sgd_optimizer.zero_grad()

    loss.backward()

    sgd_optimizer.step()



Complete Code



import torch

batch_size, input_dim, hidden_dim, out_dim = 32, 100, 100, 10

input_tensor = torch.randn(batch_size, input_dim)
output_tensor = torch.randn(batch_size, out_dim)

model = torch.nn.Sequential(
    torch.nn.Linear(input_dim, hidden_dim),
    torch.nn.Tanh(),
    torch.nn.Linear(hidden_dim, out_dim),
)


loss_function = torch.nn.MSELoss(reduction='sum')

lr = 1e-5

sgd_optimizer = torch.optim.SGD(model.parameters(), lr=lr)

for i in range(200):
    predicted_value = model(input_tensor)
    loss = loss_function(predicted_value, output_tensor)
    print(i, loss.item())

    sgd_optimizer.zero_grad()

    loss.backward()

    sgd_optimizer.step()


Output



0 314.1348571777344
1 313.5540466308594
2 312.9750061035156
3 312.3976135253906
4 311.8220520019531
5 311.2481689453125
6 310.67596435546875
7 310.10546875
8 309.53668212890625
9 308.96954345703125
10 308.40411376953125
11 307.84033203125
....

Categories: PyTorch