Views
No views yet
data_path.1import os
2
3import torch
4import torch.nn.functional as F
5from torch import nn, optim
6from torch.utils.data import DataLoader
7from torchvision.datasets import MNIST, SVHN
8from torchvision.transforms import ToTensor
9
10
11class SVHN_Classifier(nn.Module):
12 def __init__(self):
13 super(SVHN_Classifier, self).__init__()
14 self.conv1 = nn.Conv2d(3, 10, kernel_size=5)
15 self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
16 self.conv2_drop = nn.Dropout2d()
17 self.fc1 = nn.Linear(500, 50)
18 self.fc2 = nn.Linear(50, 10)
19
20 def forward(self, x):
21 x = F.relu(F.max_pool2d(self.conv1(x), 2))
22 x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
23 x = x.view(-1, 500)
24 x = F.relu(self.fc1(x))
25 x = F.dropout(x, training=self.training)
26 x = self.fc2(x)
27 return F.log_softmax(x, dim=-1)
28
29
30class MNIST_Classifier(nn.Module):
31 def __init__(self):
32 super(MNIST_Classifier, self).__init__()
33 self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
34 self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
35 self.conv2_drop = nn.Dropout2d()
36 self.fc1 = nn.Linear(320, 50)
37 self.fc2 = nn.Linear(50, 10)
38
39 def forward(self, x):
40 x = F.relu(F.max_pool2d(self.conv1(x), 2))
41 x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
42 x = x.view(-1, 320)
43 x = F.relu(self.fc1(x))
44 x = F.dropout(x, training=self.training)
45 x = self.fc2(x)
46 return F.log_softmax(x, dim=-1)
47
48
49def load_mnist_svhn_classifiers(data_path, device="cuda"):
50 c1 = MNIST_Classifier()
51 c1.load_state_dict(torch.load(f"{data_path}/mnist.pt", map_location=device))
52 c2 = SVHN_Classifier()
53 c2.load_state_dict(torch.load(f"{data_path}/svhn.pt", map_location=device))
54 return {"mnist": c1.to(device).eval(), "svhn": c2.to(device).eval()}
55