【Pytorch】nn.Module

import torch
from torch import nn

device="cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {
      device} device")

Using cuda device

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork,self).__init__()
        self.flatten=nn.Flatten()
        self.linear_relu_stack=nn.Sequential(
            nn.Linear(28*28,512),
            nn.ReLU(),
            nn.Linear(512,512),
            nn.ReLU(),
            nn.Linear(512,10),
        )
        
    def forward(self,x):
        x=self.flatten(x)
        logits=self.linear_relu_stack(x)
        return logits


model=NeuralNetwork().to(device)
model

NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)

input_image=torch.rand(3,28,28)
input_image=input_image.to("cuda")
input_image.size()

torch.Size([3, 28, 28])

logits=model(input_image)
logits.shape

torch.Size([3, 10])

softmax=nn.Softmax(dim=1)
softmax(logits)

tensor([[0.1017, 0.0993, 0.1022, 0.0997, 0.1048, 0.1015, 0.0971, 0.0979, 0.0961,
0.0996],
[0.1023, 0.1030, 0.0998, 0.1008, 0.1009, 0.1034, 0.1012, 0.0952, 0.0925,
0.1010],
[0.1048, 0.1021, 0.1008, 0.1018, 0.1034, 0.1022, 0.0991, 0.0925, 0.0965,
0.0970]], device=‘cuda:0’, grad_fn=< SoftmaxBackward >)