023/Pytorch-Project

/Pytorch-Project/Pytorch-Project/Pytorch-Project/3_10_1_1_1_1_2/3_10_1_1_1_1_2/02_model.py
import os
import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
from torch.autograd import Variable
from torch import nn
import matplotlib.pyplot as plt
import time
import random

import torchvision.models as models

def imshow(img):
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()

# 数据预处理
def data_process():
# 数据增强
data_transforms = {
'train': transforms.Compose([
transforms.RandomHorizontalFlip(),
# transforms.RandomRotation(10),
# transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
]),
'valid': transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
]),
'test': transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
}
# 数据加载
data_dir = '../datasets'
image_datasets = {x: torchvision.datasets.ImageFolder(os.path.join(data_dir, x),
transform=data_transforms[x])
for x in ['train', 'valid', 'test']}
# 数据加载器
dataloaders = {x: DataLoader(image_datasets[x], batch_size=4, shuffle=True, num_workers=4)
for x in ['train', 'valid', 'test']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'valid', 'test']}
class_names = image_datasets['train'].classes
print(dataset_sizes)
# imshow(image_datasets['train'].images[0])
print(class_names)
return image_datasets, dataloaders, dataset_sizes, class_names

class ResNet50(nn.Module):
def __init__(self):
super(ResNet50, self).__init__()
self.resnet50 = models.resnet50(pretrained=True)
self.resnet50.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.resnet50.fc = nn.Linear(2048, 2)

def forward(self, x):
x = self.resnet50(x)
return x

def main():
image_datasets, dataloaders, dataset_sizes, class_names = data_process()
resnet50 = ResNet50()
resnet50 = resnet50.cuda()
print(resnet50)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(resnet50.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4)

# 模型训练
def train_model(model, criterion, optimizer, dataloaders, dataset_sizes, num_epochs=25):
since = time.time()
best_model_wts = model.state_dict()
best_acc = 0.0

for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)

# 每个epoch有训练和验证两个阶段，每个阶段都先调用训练函数train()和验证函数test()进行训练和验证
for phase in ['train', 'valid']:
if phase == 'train':
model.train() # 训练模式
else:
model.eval() # 验证模式

running_loss = 0.0
running_corrects = 0

# 遍历数据集
for inputs, labels in dataloaders[phase]:
inputs = inputs.cuda()
labels = labels.cuda()

# 将数据和标签包装成变量
inputs = Variable(inputs)
labels = Variable(labels)

# 前向传播
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)

# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()

# 统计结果
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)

# 计算平均损失和平均准确率
epoch_loss = running_loss / dataset_sizes[phase]
epoch_acc = running_corrects.double() / dataset_sizes[phase]

print('{} Loss: {:.4f} Acc: {:.4f}'.format(
phase, epoch_loss, epoch_acc))

# 每个epoch保存模型权重
if phase == 'valid' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = model.state_dict()

print()

time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60