使用Pytorch实现Swish激活函数的示例详解

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Swish功能
class Swish(nn.Module):
	def forward(self, x):
		return x * torch.sigmoid(x)

# 定义神经网络模型
class Net(nn.Module):
	def __init__(self):
		super(Net, self).__init__()
		self.fc1 = nn.Linear(28 * 28, 256)
		self.fc2 = nn.Linear(256, 128)
		self.fc3 = nn.Linear(128, 10)
		self.swish = Swish()

	def forward(self, x):
		x = x.view(-1, 28 * 28)
		x = self.fc1(x)
		x = self.swish(x)
		x = self.fc2(x)
		x = self.swish(x)
		x = self.fc3(x)
		return x

# 创建模型的实例
model = Net()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 定义数据转换
transform = transforms.Compose([
	transforms.ToTensor(),
])

# 加载MNIST数据集
train_dataset = datasets.MNIST('', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('', train=False, download=True, transform=transform)

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 训练循环
num_epochs = 5
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model.to(device)

for epoch in range(num_epochs):
	model.train()
	total_loss = 0.0
	for batch_idx, (data, target) in enumerate(train_loader):
		data, target = data.to(device), target.to(device)
		optimizer.zero_grad()
		output = model(data)
		loss = criterion(output, target)
		loss.backward()
		optimizer.step()
		total_loss += loss.item()
	
	print(f"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss / len(train_loader)}")

# 评估循环
model.eval()
correct = 0
total = 0
with torch.no_grad():
	for data, target in test_loader:
		data, target = data.to(device), target.to(device)
		outputs = model(data)
		_, predicted = torch.max(outputs.data, 1)
		total += target.size(0)
		correct += (predicted == target).sum().item()

print(f"Accuracy on test set: {100 * correct / total}%")