DEV Community

Super Kai (Kazuya Ito)
Super Kai (Kazuya Ito)

Posted on • Edited on

Module in PyTorch

Buy Me a Coffee

*Memos:

  • My post explains Linear Regression in PyTorch.
  • My post explains Batch, Mini-Batch and Stochastic Gradient Descent with DataLoader() in PyTorch.
  • My post explains Batch Gradient Descent without DataLoader() in PyTorch.
  • My post explains how to save a model in PyTorch.
  • My post explains how to load the saved model which I show in this post in PyTorch.
  • My post explains Deep Learning Workflow in PyTorch.
  • My post explains how to clone a private repository with FGPAT(Fine-Grained Personal Access Token) from Github.
  • My post explains how to clone a private repository with PAT(Personal Access Token) from Github.
  • My post explains useful IPython magic commands.
  • My repo has models.

Module() can create a model, being its base class as shown below:

*Memos:

  • forward() must be overridden in the subclass of Module().
  • state_dict() can return the dictionary containing parameters and buffers. *It cannot get the num3 and num4 defined without Parameter().
  • load_state_dict() can load model's state_dict() into a model. *Basically, it's used to load a saved model into the currently used model.
  • parameters() can return an iterator over module parameters. *It cannot get the num3 and num4 defined without Parameter().
  • training can check if it's train mode or eval mode. By default, it's train mode.
  • train() can set a model train mode.
  • eval() can set a model eval mode.
  • cpu() can convert all model parameters and buffers to CPU.
  • cuda() can convert all model parameters and buffers to CUDA(GPU).
  • There are also save() and load(). *My post explains save() and load().
import torch
from torch import nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.num1 = nn.Parameter(torch.tensor(9.))
        self.num2 = nn.Parameter(torch.tensor(7.))
        self.num3 = torch.tensor(-2.) # Defined without `Parameter()`
        self.num4 = torch.tensor(6.) # Defined without `Parameter()`
        self.layer1 = nn.Linear(in_features=4, out_features=5)
        self.layer2 = nn.Linear(in_features=5, out_features=2)
        self.layer3 = nn.Linear(in_features=2, out_features=3)
        self.relu = nn.ReLU()

    def forward(self, x): # Must be overridden
        x1 = self.layer1(input=x)
        x2 = self.relu(input=x1)
        x3 = self.layer2(input=x2)
        x4 = self.relu(input=x3)
        x5 = self.layer3(input=x4)
        return x5

my_tensor = torch.tensor([8., -3., 0., 1.])

torch.manual_seed(42)

mymodel = MyModel()
mymodel(x=my_tensor)
# tensor([0.8092, 0.8460, 0.3758], grad_fn=<ViewBackward0>)

mymodel
# MyModel(
#   (layer1): Linear(in_features=4, out_features=5, bias=True)
#   (layer2): Linear(in_features=5, out_features=2, bias=True)
#   (layer3): Linear(in_features=2, out_features=3, bias=True)
#   (relu): ReLU()
# )

mymodel.layer2
# Linear(in_features=5, out_features=2, bias=True)

mymodel.state_dict()
# OrderedDict([('num1', tensor(9.)),
#              ('num2', tensor(7.)),
#              ('layer1.weight',
#               tensor([[0.3823, 0.4150, -0.1171, 0.4593],
#                       [-0.1096, 0.1009, -0.2434, 0.2936],
#                       [0.4408, -0.3668, 0.4346, 0.0936],
#                       [0.3694, 0.0677, 0.2411, -0.0706],
#                       [0.3854, 0.0739, -0.2334, 0.1274]])),
#              ('layer1.bias',
#               tensor([-0.2304, -0.0586, -0.2031, 0.3317, -0.3947])),
#              ('layer2.weight',
#               tensor([[-0.2062, -0.1263, -0.2689, 0.0422, -0.4417],
#                       [0.4039, -0.3799, 0.3453, 0.0744, -0.1452]])),
#              ('layer2.bias', tensor([0.2764, 0.0697])),
#              ('layer3.weight',
#               tensor([[0.5713, 0.0773],
#                       [-0.2230, 0.1900],
#                       [-0.1918, 0.2976]])),
#              ('layer3.bias', tensor([0.6313, 0.4087, -0.3091]))])

mymodel.load_state_dict(state_dict=mymodel.state_dict())
# <All keys matched successfully>

params = mymodel.parameters()

next(params)
# Parameter containing:
# tensor(9., requires_grad=True)

next(params)
# Parameter containing:
# tensor(7., requires_grad=True)

next(params)
# Parameter containing:
# tensor([[0.3823, 0.4150, -0.1171, 0.4593],
#         [-0.1096, 0.1009, -0.2434, 0.2936],
#         [0.4408, -0.3668, 0.4346, 0.0936],
#         [0.3694, 0.0677, 0.2411, -0.0706],
#         [0.3854, 0.0739, -0.2334, 0.1274]], requires_grad=True)

next(params)
# Parameter containing:
# tensor([-0.2304, -0.0586, -0.2031, 0.3317, -0.3947],
#        requires_grad=True)

next(params)
# Parameter containing:
# tensor([[-0.2062, -0.1263, -0.2689, 0.0422, -0.4417],
#         [0.4039, -0.3799, 0.3453, 0.0744, -0.1452]],
#        requires_grad=True)

next(params)
# Parameter containing:
# tensor([0.2764, 0.0697], requires_grad=True)

next(params)
# Parameter containing:
# tensor([[0.5713, 0.0773],
#         [-0.2230, 0.1900],
#         [-0.1918, 0.2976]], requires_grad=True)

next(params)
# Parameter containing:
# tensor([0.6313, 0.4087, -0.3091], requires_grad=True)

mymodel.training
# True

mymodel.eval()

mymodel.training
# False

mymodel.train()

mymodel.training
# True

mymodel.cuda(device='cuda:0')

mymodel.layer2.weight.device
# device(type='cuda', index=0)

mymodel.cpu()

mymodel.layer2.weight.device
# device(type='cpu')
Enter fullscreen mode Exit fullscreen mode

$150K MiniMax AI Agent Challenge — Build Smarter, Remix Bolder, Win Bigger!

Join the $150k MiniMax AI Agent Challenge — Build your first AI Agent 🤖

Developers, innovators, and AI tinkerers, build your AI Agent and win $150,000 in cash. 💰

Read more →

Top comments (0)

Heroku

Tired of jumping between terminals, dashboards, and code?

Check out this demo showcasing how tools like Cursor can connect to Heroku through the MCP, letting you trigger actions like deployments, scaling, or provisioning—all without leaving your editor.

Learn More

👋 Kindness is contagious

If this **helped, please leave a ❤️ or a friendly comment!

Okay