pytorch学习笔记——timm库

pip install timm

import timm

import timm

m = timm.create_model('vgg16', pretrained=True)
m.eval()

model = timm.create_model('efficientnet_b0', pretrained=True)

import timm
from pprint import pprint
model_names = timm.list_models(pretrained=True)
pprint(model_names)
>>> ['adv_inception_v3',
 'cspdarknet53',
 'cspresnext50',
 'densenet121',
 'densenet161',
 'densenet169',
 'densenet201',
 'densenetblur121d',
 'dla34',
 'dla46_c',
...
]

import timm
from pprint import pprint
model_names = timm.list_models('*resne*t*')
pprint(model_names)
>>> ['cspresnet50',
 'cspresnet50d',
 'cspresnet50w',
 'cspresnext50',
...
]

import torch.nn as nn
from timm.models.layers import ConvBnAct

block = ConvBnAct(in_channels=3, out_channels=64, kernel_size=3, stride=1, act_layer=nn.ReLU)
print(block)

ConvNormAct(
  (conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (bn): BatchNormAct2d(
    64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True
    (drop): Identity()
    (act): ReLU(inplace=True)
  )
)

import timm

class CustomModel(timm.models.VisionTransformer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.head = torch.nn.Linear(self.head.in_features, num_classes)

model = CustomModel(img_size=224, patch_size=16, embed_dim=768, depth=12, num_heads=12, num_classes=100)

import torch
import torchvision
import torchvision.transforms as transforms

# 数据预处理
transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32, padding=4),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载CIFAR-10数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=100, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)

import timm

model = timm.create_model('resnet50', pretrained=True, num_classes=10)

import torch.optim as optim

criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 10
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

for epoch in range(num_epochs):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch {epoch + 1}, Loss: {running_loss / (i + 1)}")

correct = 0
total = 0
model.eval()

with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)

        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

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

import torch
import timm
from PIL import Image
import torchvision.transforms as transforms

# 加载ResNet50模型
model = timm.create_model('vgg16', pretrained=True)

# 将模型设置为评估模式
model.eval()

# 加载图像并进行预处理
image = Image.open(r"D:\Desktop\workdata\data\segmentation_dataset\images\Abyssinian_24.jpg")
# 定义图像预处理转换
preprocess = transforms.Compose([
    transforms.Resize((256, 256)),  # 调整图像大小为256x256
    transforms.ToTensor(),  # 转换为Tensor类型
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化
])

# 应用预处理转换
processed_image = preprocess(image)

# 查看处理后的图像形状和数值范围
print("Processed image shape:", processed_image.shape)
print("Processed image range:", processed_image.min(), "-", processed_image.max())

# 5. 可选：将处理后的图像转换回PIL图像对象
processed_pil_image = transforms.ToPILImage()(processed_image)

# 6. 可选：显示处理后的图像
processed_pil_image.show()


# 将输入张量转换为批处理张量
input_batch = processed_image.unsqueeze(0)

# 将输入张量传递给模型并获取输出
with torch.no_grad():
    output = model(input_batch)

# 获取预测结果
_, predicted = torch.max(output.data, 1)

# 打印预测结果
print(predicted.item())

import torch
import timm
m = timm.create_model('resnet50', pretrained=True, num_classes=1000)
o = m(torch.randn(2, 3, 224, 224))
print(f'Pooled shape: {o.shape}')
# Pooled shape: torch.Size([2, 2048])

import torch
import timm
m = timm.create_model('ese_vovnet19b_dw', pretrained=True)
o = m(torch.randn(2, 3, 224, 224))
print(f'Original shape: {o.shape}')
m.reset_classifier(0)
o = m(torch.randn(2, 3, 224, 224))
print(f'Pooled shape: {o.shape}')
# Pooled shape: torch.Size([2, 1024])

import torch
import timm
m = timm.create_model('resnest26d', features_only=True, pretrained=True)
o = m(torch.randn(2, 3, 224, 224))
for x in o:
  print(x.shape)

# output
torch.Size([2, 64, 112, 112])
torch.Size([2, 256, 56, 56])
torch.Size([2, 512, 28, 28])
torch.Size([2, 1024, 14, 14])
torch.Size([2, 2048, 7, 7])

import torch
import timm

# 1. 加载预训练模型
model = timm.create_model('resnet50', pretrained=True)

# 2. 移除最后一层分类器
model = torch.nn.Sequential(*list(model.children())[:-1])

# 3. 设置模型为评估模式
model.eval()

# 4. 加载图像
input = torch.randn(1, 3, 224, 224)  # 假设输入为224x224的RGB图像

# 5. 特征提取
features = model(input)

# 6. 打印提取到的特征
print("Features shape:", features.shape)

import torch
import timm

class ModelEnsemble(torch.nn.Module):
    def __init__(self, models):
        super().__init__()
        self.models = torch.nn.ModuleList(models)

    def forward(self, x):
        outputs = [model(x) for model in self.models]
        return torch.mean(torch.stack(outputs), dim=0)

# 加载多个预训练模型
model1 = timm.create_model('resnet18', pretrained=True, num_classes=100)
model2 = timm.create_model('vgg16', pretrained=True, num_classes=100)

# 创建模型融合
ensemble = ModelEnsemble([model1, model2])

# 使用融合模型进行预测
output = ensemble(input_tensor)

ModelEnsemble(
  (models): ModuleList(
    (0): ResNet(
      (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (act1): ReLU(inplace=True)
      (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
      (layer1): Sequential(
      ........ （这里省略了ResNet的网络结构）
      )
      (global_pool): SelectAdaptivePool2d (pool_type=avg, flatten=Flatten(start_dim=1, end_dim=-1))
      (fc): Linear(in_features=512, out_features=100, bias=True)
    )
    (1): VGG(
      (features): Sequential(
        (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): ReLU(inplace=True)
        (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (3): ReLU(inplace=True)
        (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        .......(这里省略了VGG16的模型结构)
      )
      (pre_logits): ConvMlp(
        (fc1): Conv2d(512, 4096, kernel_size=(7, 7), stride=(1, 1))
        (act1): ReLU(inplace=True)
        (drop): Dropout(p=0.0, inplace=False)
        (fc2): Conv2d(4096, 4096, kernel_size=(1, 1), stride=(1, 1))
        (act2): ReLU(inplace=True)
      )
      (head): ClassifierHead(
        (global_pool): SelectAdaptivePool2d (pool_type=avg, flatten=Flatten(start_dim=1, end_dim=-1))
        (fc): Linear(in_features=4096, out_features=100, bias=True)
        (flatten): Identity()
      )
    )
  )
)

import torch
import timm

# 1. 加载并初始化要融合的模型
model1 = timm.create_model('resnet50', pretrained=True)
model2 = timm.create_model('efficientnet_b0', pretrained=True)
model3 = timm.create_model('densenet121', pretrained=True)

# 2. 定义融合模型
class FusionModel(torch.nn.Module):
    def __init__(self, model1, model2, model3):
        super(FusionModel, self).__init__()
        self.model1 = model1
        self.model2 = model2
        self.model3 = model3
        self.fc = torch.nn.Linear(3, 1)  # 假设最终输出为单个值

    def forward(self, x):
        output1 = self.model1(x)
        output2 = self.model2(x)
        output3 = self.model3(x)
        fused_output = torch.cat([output1, output2, output3], dim=1)
        fused_output = self.fc(fused_output)
        return fused_output

# 3. 创建融合模型实例
fusion_model = FusionModel(model1, model2, model3)

# 4. 使用融合模型进行推理
input = torch.randn(1, 3, 224, 224)  # 假设输入为224x224的RGB图像
output = fusion_model(input)

# 5. 打印输出结果
print("Fused output shape:", output.shape)