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 >)