请帮我写出基于pytorch的ResNet34网络实现图像二分类的代码并且包含训练集、验证集和测试集代码,且用tqdm代码实现训练进度
时间: 2024-02-27 10:58:26 浏览: 45
好的,以下是基于pytorch的ResNet34网络实现图像二分类的代码,包括训练集、验证集和测试集代码,并使用tqdm实现训练进度。注意,代码中的路径需要根据实际数据集路径进行修改。
```python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from tqdm import tqdm
# 定义训练、验证、测试数据集路径
train_data_path = "path/to/train/dataset"
val_data_path = "path/to/val/dataset"
test_data_path = "path/to/test/dataset"
# 定义训练、验证、测试数据集的transform
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# 加载数据集
train_dataset = torchvision.datasets.CIFAR10(root=train_data_path, train=True, download=True, transform=transform_train)
val_dataset = torchvision.datasets.CIFAR10(root=val_data_path, train=False, download=True, transform=transform_test)
test_dataset = torchvision.datasets.CIFAR10(root=test_data_path, train=False, download=True, transform=transform_test)
# 定义超参数
batch_size = 64
learning_rate = 0.001
num_epochs = 10
# 定义数据加载器
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# 定义ResNet34模型
class ResNet34(nn.Module):
def __init__(self, num_classes):
super(ResNet34, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.layer1 = self._make_layer(64, 3)
self.layer2 = self._make_layer(128, 4, stride=2)
self.layer3 = self._make_layer(256, 6, stride=2)
self.layer4 = self._make_layer(512, 3, stride=2)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512, num_classes)
def _make_layer(self, planes, blocks, stride=1):
downsample = None
if stride != 1 or planes != 64:
downsample = nn.Sequential(
nn.Conv2d(64, planes, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes)
)
layers = []
layers.append(BasicBlock(64, planes, stride, downsample))
for i in range(1, blocks):
layers.append(BasicBlock(planes, planes))
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
class BasicBlock(nn.Module):
def __init__(self, in_planes, planes, stride=1, downsample=None):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.downsample = downsample
def forward(self, x):
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out
# 定义模型、损失函数和优化器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ResNet34(num_classes=2).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# 训练模型
for epoch in range(num_epochs):
train_loss = 0.0
train_correct = 0
train_total = 0
model.train()
for images, labels in tqdm(train_loader, desc=f"Epoch {epoch + 1}/{num_epochs}"):
images = images.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
train_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
train_total += labels.size(0)
train_correct += (predicted == labels).sum().item()
train_accuracy = 100 * train_correct / train_total
train_loss /= len(train_loader)
val_loss = 0.0
val_correct = 0
val_total = 0
model.eval()
with torch.no_grad():
for images, labels in val_loader:
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
val_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
val_total += labels.size(0)
val_correct += (predicted == labels).sum().item()
val_accuracy = 100 * val_correct / val_total
val_loss /= len(val_loader)
print(f"Epoch {epoch + 1}/{num_epochs}, Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.2f}%, Val Loss: {val_loss:.4f}, Val Accuracy: {val_accuracy:.2f}%")
# 测试模型
test_loss = 0.0
test_correct = 0
test_total = 0
model.eval()
with torch.no_grad():
for images, labels in tqdm(test_loader, desc="Testing"):
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
test_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
test_total += labels.size(0)
test_correct += (predicted == labels).sum().item()
test_accuracy = 100 * test_correct / test_total
test_loss /= len(test_loader)
print(f"Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.2f}%")
```
阅读全文