Views
No views yet
1import timm
2import torchvision
3MNIST_PATH = './datasets/mnist'
4
5
6net = timm.create_model("resnet18", pretrained=False, num_classes=10)
7net.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
8net.load_state_dict(
9 torch.hub.load_state_dict_from_url(
10 "https://huggingface.co/gpcarl123/resnet18_mnist/resolve/main/resnet18_mnist.pth",
11 map_location="cpu",
12 file_name="resnet18_mnist.pth",
13 )
14)
15
16preprocessor = torchvision.transforms.Normalize((0.1307,), (0.3081,))
17transform = transforms.Compose([transforms.ToTensor()])
18test_set = datasets.MNIST(root=MNIST_PATH, train=False, download=True, transform=transform)
19test_loader = data.DataLoader(test_set, batch_size=5, shuffle=False, num_workers=2)
20
21for data, target in test_loader:
22 print(net(preprocessor(data)))
23 print(target)
24 break
25