Swin Transformer猫狗分类
前言
早上看了一下Swin Transformer的论文，觉得还不错，就看了看代码，还挺简洁。
我不说是谁，那么无聊画了一下午用Swin Tranformer实现猫狗分类…
代码
依赖
需要下载一个库, 在终端运行则不需要前面的英文感叹号
!pip install timm
将依赖import 进来
import torch
import torchvision
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from torchvision import models
from timm.models.swin_transformer import SwinTransformer
import numpy as np
import os
import cv2
import random
import shutil
from tqdm import tqdm
设置一些超参数
batch_size_train = 32
batch_size_test = 32
learning_rate = 0.01
momentum = 0.5
log_interval = 10
random_seed = 1
torch.manual_seed(random_seed)
如果有gpu则使用gpu加速
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(device)
数据准备
下载数据
!wget https://download.microsoft.com/download/3/E/1/3E1C3F21-ECDB-4869-8368-6DEBA77B919F/kagglecatsanddogs_3367a.zip
查看一下大小
!du -sh /content/kagglecatsanddogs_3367a.zip
# 787M
解压
!unzip kagglecatsanddogs_3367a.zip
显示一下部分图片，其实图片是彩色的，只是这段代码写的不那么好，出来时灰色的，以后再改。另外从图片看出，图片不是正方形，不是224224 或 384384，所以后面读取的时候需要resize一下。
# 本小块展示的代码来自: https://blog.csdn.net/Neuf_Soleil/article/details/83957756
DATADIR = "/content/PetImages"
CATEGORIES = ["Dog", "Cat"]
for category in CATEGORIES:
path = os.path.join(DATADIR,category)
for img in os.listdir(path):
img_array = cv2.imread(os.path.join(path, img) ,cv2.IMREAD_GRAYSCALE) # 把图片读取为数组
plt.imshow(img_array, cmap='gray') # 使用灰度图
plt.show()
break # 这里先拿一个图片测试，所以 break 两次
划分训练集和测试集
注意，有时候你的数据集已经划分好了，所以不需要该步骤。但是这里我们下载的数据集是这样的:
| PetImages
| --- Cat
| --- Dog
cat_file_lt = list(os.listdir('/content/PetImages/Cat'))
dog_file_lt = list(os.listdir('/content/PetImages/Dog'))
cat_file_total_num = len(cat_file_lt)
dog_file_total_num = len(dog_file_lt)
print('cat: ', cat_file_total_num)
print('dog: ', dog_file_total_num)
dir_lt = ['data', 'data/train', 'data/test',
'data/train/Cat', 'data/train/Dog',
'data/test/Cat', 'data/test/Dog']
for dir in dir_lt:
if not os.path.exists(dir):
os.mkdir(dir)
# 打乱数据
random.shuffle(cat_file_lt)
random.shuffle(dog_file_lt)
def copy_from_lt1_to_lt2(origin_dir, origin_lt, target_dir):
for pic in tqdm(origin_lt):
origin_path = os.path.join(origin_dir, pic)
target_path = os.path.join(target_dir, pic)
shutil.copyfile(origin_path, target_path)
# 做个例子，拿1000张...
# train_test_split = int(cat_file_total_num * 0.8)
train_cat_lt = cat_file_lt[:1000]
train_dog_lt = dog_file_lt[:1000]
test_cat_lt = cat_file_lt[1000:1200]
test_dog_lt = dog_file_lt[1000:1200]
train_test_path_lt = [('/content/PetImages/Cat', train_cat_lt, '/content/data/train/Cat'),
('/content/PetImages/Dog', train_dog_lt, '/content/data/train/Dog'),
('/content/PetImages/Cat', test_cat_lt, '/content/data/test/Cat'),
('/content/PetImages/Dog', test_dog_lt, '/content/data/test/Dog')]
copy_from_lt1_to_lt2(*(train_test_path_lt[0]))
copy_from_lt1_to_lt2(*(train_test_path_lt[1]))
copy_from_lt1_to_lt2(*(train_test_path_lt[2]))
copy_from_lt1_to_lt2(*(train_test_path_lt[3]))
处理完后目录结构如下:
读取图片
transform = torchvision.transforms.Compose(
[
torchvision.transforms.Resize((224, 224)), # 重置图片大小
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_dataset = torchvision.datasets.ImageFolder(root='/content/data/train', transform=transform)
test_dataset = torchvision.datasets.ImageFolder(root='/content/data/test', transform=transform)
print('len(train_dataset): ', len(train_dataset))
print('len(test_dataset): ', len(test_dataset))
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size_train, shuffle=True, num_workers=2)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size_test, shuffle=True, num_workers=2)
print('len(train_loader): ', len(train_loader))
print('len(test_loader): ', len(test_loader))
# 查看一个批次的数据
examples = enumerate(train_loader)
batch_idx, (example_data, example_targets) = next(examples)
print(example_targets)
print(example_data.shape) # [32, 3, 224, 224]
# 本小块展示的代码来自 https://blog.csdn.net/theVicTory/article/details/109230519
import matplotlib.pyplot as plt
import numpy as np
# 输出图像的函数
def imshow(img):
img = img / 2 + 0.5 # 反标准化
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()
# 获取一个批次的训练图片、标签并显示
images, labels = iter(train_loader).next()
imshow(torchvision.utils.make_grid(images))
模型定义
模型定义可参考原论文的表格。
这里我用Swin transformer的Tiny版本，它的参数量和ResNet50差不多。
swin_tiny_cfg = dict(patch_size=4, window_size=7, embed_dim=96, depths=(2, 2, 6, 2), num_heads=(3, 6, 12, 24))
swin_tiny = SwinTransformer(**swin_tiny_cfg)
下载权重
!wget https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.backbone = SwinTransformer(**swin_tiny_cfg)
self.backbone.load_state_dict(
torch.load('/content/swin_tiny_patch4_window7_224.pth',
map_location=torch.device(device))['model'], strict=True)
self.fc1 = nn.Linear(768, 256)
self.fc2 = nn.Linear(256, 2)
self.relu = nn.ReLU()
def forward(self, x):
"""
Swin Transformer https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/swin_transformer.py
def forward_features(self, x):
x = self.patch_embed(x)
if self.absolute_pos_embed is not None:
x = x + self.absolute_pos_embed
x = self.pos_drop(x)
x = self.layers(x)
x = self.norm(x) # B L C
x = self.avgpool(x.transpose(1, 2)) # B C 1
x = torch.flatten(x, 1)
return x
def forward(self, x):
x = self.forward_features(x)
x = self.head(x)
return x
"""
x = self.backbone.forward_features(x) # [batch_size, 768]
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x) # [batch_size, 2]
return x
model = Net()
# output = model(example_data)
# print(output.shape)
model = model.to(device)
看看参数量
def get_parameter_number(model_analyse):
# 打印模型参数量
total_num = sum(p.numel() for p in model_analyse.parameters())
trainable_num = sum(p.numel() for p in model_analyse.parameters() if p.requires_grad)
return 'Total parameters: {}, Trainable parameters: {}'.format(total_num, trainable_num)
# 查看一下模型总的参数量和可学习参数量
get_parameter_number(model)
训练与预测
损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate,
momentum=momentum)
记录训练过程的损失变化
train_losses = []
train_counter = []
test_losses = []
def train(epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
data = data.to(device)
target = target.to(device)
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
if batch_idx % log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
train_losses.append(loss.item())
train_counter.append(
(batch_idx * batch_size_train) + ((epoch - 1) * len(train_loader.dataset)))
torch.save(model.state_dict(), './model.pth')
torch.save(optimizer.state_dict(), './optimizer.pth')
作为演示，只训练一个epoch，通常是几个。
train(epoch=1)
def test():
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data = data.to(device)
target = target.to(device)
output = model(data)
test_loss += criterion(output, target).item()
pred = output.data.max(1, keepdim=True)[1]
correct += pred.eq(target.data.view_as(pred)).sum()
test_loss /= len(test_loader.dataset)
test_losses.append(test_loss)
print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
这准确率很惊人…虽然只是二分类问题…
test()
画出训练时损失下降趋势
fig = plt.figure()
plt.plot(train_counter, train_losses, color='blue')
plt.legend(['Train Loss'], loc='upper right')
plt.xlabel('number of training examples seen')
plt.ylabel('negative log likelihood loss')
plt.show()
