Views
No views yet
MyNeuralNet is a simple, fully connected neural network designed for classifying the handwritten digits of the MNIST dataset. The model consists of three linear layers with ReLU activation functions, followed by a final layer with a softmax output to predict probabilities across the 10 possible digits (0-9).1for epoch in range(n_epochs):
2 for images, labels in dataloader:
3 # Forward pass
4 predictions = model(images)
5 loss = loss_function(predictions, labels)
6
7 # Backward pass and optimization
8 optimizer.zero_grad()
9 loss.backward()
10 optimizer.step()MyNeuralNet and use it to predict MNIST images:1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4from torch.utils.data import Dataset, DataLoader
5from huggingface_hub import hf_hub_download
6
7
8# Ensure the device selection logic is centralized
9def get_device():
10 return torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
12
13# Define the neural network architecture
14class MyNeuralNet(nn.Module):
15 def __init__(self):
16 super(MyNeuralNet, self).__init__()
17 self.Matrix1 = nn.Linear(28 * 28, 100)
18 self.Matrix2 = nn.Linear(100, 50)
19 self.Matrix3 = nn.Linear(50, 10)
20 self.R = nn.ReLU()
21
22 def forward(self, x):
23 x = x.view(-1, 28 * 28)
24 x = self.R(self.Matrix1(x))
25 x = self.R(self.Matrix2(x))
26 x = self.Matrix3(x)
27 return x.squeeze()
28
29
30# Define the custom dataset class
31class CTDataset(Dataset):
32 def __init__(self, filepath, device):
33 # Add 'device' as a parameter to the class constructor
34 x, y = torch.load(filepath)
35 self.x = x.float().div(255).to(device) # Use the passed 'device' for tensor operations
36 self.y = F.one_hot(y, num_classes=10).float().to(device)
37
38 def __len__(self):
39 return self.x.shape[0]
40
41 def __getitem__(self, ix):
42 return self.x[ix], self.y[ix]
43
44
45def load_model():
46 device = get_device()
47 model_state_dict = torch.load(hf_hub_download(repo_id="Svenni551/may-mnist-digits", filename="model.pth"),
48 map_location=torch.device(device))
49 model = MyNeuralNet().to(device)
50 model.load_state_dict(model_state_dict)
51 model.eval()
52 return model
53
54
55def predict(input_data):
56 device = get_device()
57 model = load_model()
58 if isinstance(input_data, str): # Assuming filepath to dataset
59 dataset = CTDataset(input_data, device) # Pass 'device' as an argument
60 loader = DataLoader(dataset, batch_size=32, shuffle=False)
61 predictions = []
62 with torch.no_grad():
63 for batch, _ in loader:
64 yhat = model(batch).argmax(axis=1).cpu().numpy()
65 predictions.extend(yhat)
66 return predictions
67 elif isinstance(input_data, torch.Tensor):
68 if len(input_data.shape) == 3: # Single image
69 input_data = input_data.unsqueeze(0) # Add batch dimension
70 input_data = input_data.to(device)
71 with torch.no_grad():
72 prediction = model(input_data).argmax(axis=1).item()
73 return prediction
74 else:
75 raise ValueError("Unsupported input type. Provide a file path to a dataset or a PyTorch Tensor.")
76
77# Example usage:
78# prediction = predict('path/to/your/dataset.pt')
79# or for an image:
80# prediction = predict(your_image_tensor)
81
82# print(prediction)



