131 lines
3.6 KiB
Python
131 lines
3.6 KiB
Python
import torch
|
|
import torchvision
|
|
import torchvision.transforms as transforms
|
|
from torch import optim
|
|
from torch.utils.data import Dataset
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
data_path = "data/"
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.mps.is_available() else "cpu")
|
|
|
|
transforms = transforms.Compose([
|
|
transforms.ToTensor(), # 將圖像轉換為 Tensor
|
|
transforms.Normalize((0.4915, 0.4823, 0.4468), (0.2470, 0.2435, 0.2616))
|
|
# 歸一化，第一個 tuple 代表 CIFAR-10 這個資料集 RGB 三個通道的平均值，第二個 tuple 代表標準差
|
|
])
|
|
|
|
train_dataset = torchvision.datasets.CIFAR10(
|
|
root=data_path,
|
|
train=True,
|
|
download=True,
|
|
transform=transforms
|
|
)
|
|
|
|
test_dataset = torchvision.datasets.CIFAR10(
|
|
root=data_path,
|
|
train=False,
|
|
download=True,
|
|
transform=transforms
|
|
)
|
|
|
|
label_map = {
|
|
0: 0, # 飛機
|
|
2: 1, # 小鳥
|
|
}
|
|
|
|
class_names = ["airplane", "bird"]
|
|
|
|
train_dataset = [(img, label_map[label]) for img, label in train_dataset if label in [0, 2]]
|
|
test_dataset = [(img, label_map[label]) for img, label in test_dataset if label in [0, 2]]
|
|
|
|
|
|
class ModelDataset(Dataset):
|
|
def __init__(self, dataset):
|
|
self.dataset = dataset
|
|
|
|
def __getitem__(self, index):
|
|
img, label = self.dataset[index]
|
|
return img, label
|
|
|
|
def __len__(self):
|
|
return len(self.dataset)
|
|
|
|
|
|
train_dataset = ModelDataset(train_dataset)
|
|
test_dataset = ModelDataset(test_dataset)
|
|
|
|
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
|
|
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=True)
|
|
|
|
# 宣告模型
|
|
class Net(nn.Module):
|
|
def __init__(self, n_chansl=16):
|
|
super().__init__()
|
|
|
|
self.n_chansl = n_chansl
|
|
|
|
self.conv1 = nn.Conv2d(3, n_chansl, kernel_size=3, padding=1)
|
|
self.conv1_batchnorm = nn.BatchNorm2d(n_chansl) # 歸一化器，輸入為通道數，輸出為相同的通道數
|
|
self.conv2 = nn.Conv2d(n_chansl, n_chansl // 2, kernel_size=3, padding=1)
|
|
self.conv2_batchnorm = nn.BatchNorm2d(n_chansl // 2)
|
|
|
|
self.fc1 = nn.Linear(8 * 8 * n_chansl // 2, 32)
|
|
self.fc2 = nn.Linear(32, 2)
|
|
|
|
def forward(self, x):
|
|
out = F.max_pool2d(self.conv1_batchnorm(torch.relu(self.conv1(x))), kernel_size=2)
|
|
out = F.max_pool2d(self.conv2_batchnorm(torch.relu(self.conv2(out))), kernel_size=2)
|
|
out = out.view(-1, 8 * 8 * self.n_chansl // 2)
|
|
out = torch.relu(self.fc1(out))
|
|
out = self.fc2(out)
|
|
return out
|
|
|
|
model = Net().to(device)
|
|
|
|
optimizer = optim.SGD(model.parameters(), lr=1e-2)
|
|
loss_fn = nn.CrossEntropyLoss()
|
|
|
|
def train(epoch):
|
|
global loss
|
|
for epoch in range(epoch):
|
|
for (image, label) in train_loader:
|
|
image = image.to(device)
|
|
label = label.to(device)
|
|
outputs = model(image)
|
|
loss = loss_fn(outputs, label)
|
|
|
|
l2_lambda = 0.001
|
|
l2_norm = sum(p.pow(2.0).sum() for p in model.parameters())
|
|
loss = loss + l2_lambda * l2_norm
|
|
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
print(f"Epoch {epoch}, Loss {loss}")
|
|
|
|
|
|
train(100)
|
|
|
|
|
|
def test():
|
|
correct = 0
|
|
total = 0
|
|
with torch.no_grad():
|
|
for (image, label) in test_loader:
|
|
image = image.to(device)
|
|
label = label.to(device)
|
|
outputs = model(image)
|
|
_, predicted = torch.max(outputs, dim=1)
|
|
total += label.size(0)
|
|
correct += (predicted == label).sum().item()
|
|
|
|
print(f"Accuracy: {correct / total}")
|
|
|
|
|
|
test()
|
|
|
|
# torch.save(model.state_dict(), "model/model.pt")
|