1. 首先需要准备数据集,数据集模式按照如下格式进行准备,测试集也一样按照如下格式准备。
Data
----train
----mask
----nomask
----val
----mask
----nomask
2. 数据集加载和预处理
import torchvision.transforms as transforms
transform = transforms.Compose([
transforms.Resize((224,224)),
transforms.ToTensor(),
transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))
])
2. 分别写入训练集跟验证集
# 分别写入训练集和测试集的路径哦
dataset_train = datasets.ImageFolder('/content/Data/train/', transform)
print(dataset_train.imgs)
print(dataset_train.class_to_idx)
dataset_test = datasets.ImageFolder('/content/Data/val/', transform_test)
print(dataset_test.class_to_idx)
# 获取数据
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)
#学习率
modellr = 1e-4
#损失函数
criterion = nn.CrossEntropyLoss()
#model(resnet18)创建一个新的模型,使用resnet进行分类
model = torchvision.models.resnet50(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)
model.to(DEVICE)
#优化器,通常是使用adam
optimizer = optim.Adam(model.parameters(), lr=modellr)
3. 分别进行训练和验证
def adjust_learning_rate(optimizer, epoch):
"""Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
modellrnew = modellr * (0.1 ** (epoch // 50))
print("lr:", modellrnew)
for param_group in optimizer.param_groups:
param_group['lr'] = modellrnew
# trainning
def train(model, device, train_loader, optimizer, epoch):
model.train()
sum_loss = 0
total_num = len(train_loader.dataset)
print(total_num, len(train_loader))
for batch_idx, (data, target) in enumerate(train_loader):
data, target = Variable(data).to(device), Variable(target).to(device)
output = model(data)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print_loss = loss.data.item()
sum_loss += print_loss
if (batch_idx + 1) % 50 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
100. * (batch_idx + 1) / len(train_loader), loss.item()))
ave_loss = sum_loss / len(train_loader)
print('epoch:{},loss:{}'.format(epoch, ave_loss))
# cross-validation
def val(model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
total_num = len(test_loader.dataset)
print(total_num, len(test_loader))
with torch.no_grad():
for data, target in test_loader:
data, target = Variable(data).to(device), Variable(target).to(device)
output = model(data)
loss = criterion(output, target)
_, pred = torch.max(output.data, 1)
correct += torch.sum(pred == target)
print_loss = loss.data.item()
test_loss += print_loss
correct = correct.data.item()
acc = correct / total_num
avgloss = test_loss / len(test_loader)
print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
avgloss, correct, len(test_loader.dataset), 100 * acc))
temp_acc = 100 * acc
# saving model
for epoch in range(1, EPOCHS + 1):
adjust_learning_rate(optimizer, epoch)
train(model, DEVICE, train_loader, optimizer, epoch)
val(model, DEVICE, test_loader)
torch.save(model, 'model' + str(epoch) + '.pth')
测试结果:
# test_dataset loading
classes = ('mask', 'nomask')
transform_test = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torch.load("model4.pth")
model.eval()
model.to(DEVICE)
delDir3 = "/content/Data/test/mask/.ipynb_checkpoints"
if(os.path.exists(delDir3)):
os.rmdir(delDir3)
path='/content/Data/test/mask/'
testList=os.listdir(path)
# begin testing
count = 0
total = 0
for file in testList:
img=Image.open(path+file)
img=transform_test(img)
img.unsqueeze_(0)
img = Variable(img).to(DEVICE)
out=model(img)
# Predict
_, pred = torch.max(out.data, 1)
total = total + 1
if pred.data.item() == 0:
count = count + 1
print('Image Name:{},predict:{},predict_label:{}'.format(file,pred.data.item(),classes[pred.data.item()]))
print('mask data test accuracy:{}'.format(count/total))
count1 = 0
total1 = 0
delDir4 = "/content/Data/test/nomask/.ipynb_checkpoints"
if(os.path.exists(delDir4)):
os.rmdir(delDir4)
path='/content/Data/test/nomask/'
testList=os.listdir(path)
# begin testing
for file in testList:
img=Image.open(path+file)
img=transform_test(img)
img.unsqueeze_(0)
img = Variable(img).to(DEVICE)
out=model(img)
# Predict
_, pred = torch.max(out.data, 1)
total1 = total1 + 1
if pred.data.item() == 1:
count1 = count1 + 1
print('Image Name:{},predict:{},predict_label:{}'.format(file,pred.data.item(),classes[pred.data.item()]))
print('nomask data test accuracy:{}'.format(count1/total1))
def g_mean(tp,fn,fp,tn):
tpr = tp/(tp+fn)
tnr = tn/(tn+fp)
ret = math.sqrt(tpr*tnr)
return ret
value = g_mean(count,total-count,total1-count1,count1)
print('G-mean is :'+str(value))
想看可视化显示的看这篇文章:不多赘述了
Pytorch 实现Resnet分类Cifar10数据集>95%_resnet cifar-10 pytorch-CSDN博客