Views
No views yet
1import torch
2import torch.nn as nn
3
4class MLP(nn.Module):
5 def __init__(self, hidden_size=128, dropout=0.2):
6 super().__init__()
7 self.net = nn.Sequential(
8 nn.Flatten(),
9 nn.Linear(28*28, hidden_size), nn.ReLU(), nn.Dropout(dropout),
10 nn.Linear(hidden_size, hidden_size), nn.ReLU(), nn.Dropout(dropout),
11 nn.Linear(hidden_size, 10),
12 )
13 def forward(self, x):
14 return self.net(x)
15
16from huggingface_hub import hf_hub_download
17path = hf_hub_download(repo_id="你的用户名/mnist-mlp", filename="model.pth")
18model = MLP()
19model.load_state_dict(torch.load(path, map_location="cpu"))
20model.eval()