85 lines
2.7 KiB
Python
85 lines
2.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torch.utils.data import DataLoader, TensorDataset
|
|
|
|
|
|
class CNN(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = nn.Sequential(
|
|
nn.Conv2d(1, 32, kernel_size=3), # 28x28 -> 26x26
|
|
nn.ReLU(),
|
|
nn.MaxPool2d(2), # 26x26 -> 13x13
|
|
nn.Conv2d(32, 64, kernel_size=3), # 13x13 -> 11x11
|
|
nn.ReLU(),
|
|
nn.MaxPool2d(2), # 11x11 -> 5x5
|
|
)
|
|
self.fc = nn.Sequential(
|
|
nn.Flatten(),
|
|
nn.Linear(64 * 5 * 5, 128),
|
|
nn.ReLU(),
|
|
nn.Linear(128, 10),
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = x.view(-1, 1, 28, 28)
|
|
return self.fc(self.conv(x))
|
|
|
|
|
|
class CNNClassifier:
|
|
def __init__(self, epochs=5):
|
|
self.epochs = epochs
|
|
|
|
def fit(self, X, y):
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
self._device = device
|
|
|
|
X_tr = torch.tensor(X, dtype=torch.float32)
|
|
y_tr = torch.tensor(y, dtype=torch.long)
|
|
|
|
# Hold out 10% of the training data to track progress each epoch
|
|
n_val = len(X_tr) // 10
|
|
X_val, X_tr = X_tr[:n_val], X_tr[n_val:]
|
|
y_val, y_tr = y_tr[:n_val], y_tr[n_val:]
|
|
|
|
loader = DataLoader(TensorDataset(X_tr, y_tr), batch_size=64, shuffle=True)
|
|
|
|
model = CNN().to(device)
|
|
optimizer = optim.Adam(model.parameters(), lr=1e-3)
|
|
loss_fn = nn.CrossEntropyLoss()
|
|
|
|
print(f"\nTraining CNN (epochs={self.epochs})")
|
|
for epoch in range(1, self.epochs + 1):
|
|
model.train()
|
|
total_loss = 0
|
|
for xb, yb in loader:
|
|
xb, yb = xb.to(device), yb.to(device)
|
|
optimizer.zero_grad()
|
|
loss = loss_fn(model(xb), yb)
|
|
loss.backward()
|
|
optimizer.step()
|
|
total_loss += loss.item()
|
|
|
|
model.eval()
|
|
with torch.no_grad():
|
|
val_pred = model(X_val.to(device)).argmax(dim=1).cpu()
|
|
val_accuracy = (val_pred == y_val).float().mean().item()
|
|
|
|
print(f" epoch {epoch:2d}/{self.epochs} loss={total_loss / len(loader):.3f} val_accuracy={val_accuracy:.3f}")
|
|
print()
|
|
|
|
self._model = model
|
|
return self
|
|
|
|
def predict_proba(self, X):
|
|
X_te = torch.tensor(X, dtype=torch.float32)
|
|
self._model.eval()
|
|
with torch.no_grad():
|
|
logits = self._model(X_te.to(self._device))
|
|
probabilities = torch.softmax(logits, dim=1).cpu().numpy()
|
|
return probabilities
|
|
|
|
def predict(self, X):
|
|
return self.predict_proba(X).argmax(axis=1)
|