ViT Vision Transformer超详细解析,网络构建,可视化,数据预处理,全流程实例教程

news2024/11/8 15:30:38

关于ViT的分析和教程,网上又虚又空的东西比较多,本文通过一个实例,将ViT全解析。

包括三部分内容,网络构建;orchview.draw_graph 将网络每一层的结构与输入输出可视化;数据预处理。附完整代码

网络构建

创建一个model.py,其中实现ViT网络构建

import torch.nn as nn
import torch
import torch.optim as optim
import torch.nn.functional as F
import lightning as L


class AttentionBlock(nn.Module):
    def __init__(self, embed_dim, hidden_dim, num_heads, dropout=0.0):
        """
        Inputs:
            embed_dim - Dimensionality of input and attention feature vectors
            hidden_dim - Dimensionality of hidden layer in feed-forward network
                         (usually 2-4x larger than embed_dim)
            num_heads - Number of heads to use in the Multi-Head Attention block
            dropout - Amount of dropout to apply in the feed-forward network
        """
        super().__init__()

        self.layer_norm_1 = nn.LayerNorm(embed_dim)
        self.attn = nn.MultiheadAttention(embed_dim, num_heads)
        self.layer_norm_2 = nn.LayerNorm(embed_dim)
        self.linear = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, embed_dim),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        inp_x = self.layer_norm_1(x)
        x = x + self.attn(inp_x, inp_x, inp_x)[0]
        x = x + self.linear(self.layer_norm_2(x))
        return x


class VisionTransformer(nn.Module):
    def __init__(
        self,
        embed_dim,
        hidden_dim,
        num_channels,
        num_heads,
        num_layers,
        num_classes,
        patch_size,
        num_patches,
        dropout=0.0,
    ):
        """
        Inputs:
            embed_dim - Dimensionality of the input feature vectors to the Transformer
            hidden_dim - Dimensionality of the hidden layer in the feed-forward networks
                         within the Transformer
            num_channels - Number of channels of the input (3 for RGB)
            num_heads - Number of heads to use in the Multi-Head Attention block
            num_layers - Number of layers to use in the Transformer
            num_classes - Number of classes to predict
            patch_size - Number of pixels that the patches have per dimension
            num_patches - Maximum number of patches an image can have
            dropout - Amount of dropout to apply in the feed-forward network and
                      on the input encoding
        """
        super().__init__()

        self.patch_size = patch_size

        # Layers/Networks
        self.input_layer = nn.Linear(num_channels * (patch_size**2), embed_dim)
        self.transformer = nn.Sequential(
            *(AttentionBlock(embed_dim, hidden_dim, num_heads, dropout=dropout) for _ in range(num_layers))
        )
        self.mlp_head = nn.Sequential(nn.LayerNorm(embed_dim), nn.Linear(embed_dim, num_classes))
        self.dropout = nn.Dropout(dropout)

        # Parameters/Embeddings
        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))
        self.pos_embedding = nn.Parameter(torch.randn(1, 1 + num_patches, embed_dim))

    def img_to_patch(self, x, patch_size, flatten_channels=True):
        """
        Inputs:
            x - Tensor representing the image of shape [B, C, H, W]
            patch_size - Number of pixels per dimension of the patches (integer)
            flatten_channels - If True, the patches will be returned in a flattened format
                               as a feature vector instead of a image grid.
        """
        B, C, H, W = x.shape
        x = x.reshape(B, C, H // patch_size, patch_size, W // patch_size, patch_size)
        x = x.permute(0, 2, 4, 1, 3, 5)  # [B, H', W', C, p_H, p_W]
        x = x.flatten(1, 2)  # [B, H'*W', C, p_H, p_W]
        if flatten_channels:
            x = x.flatten(2, 4)  # [B, H'*W', C*p_H*p_W]
        return x

    def forward(self, x):
        # Preprocess input
        x = self.img_to_patch(x, self.patch_size)
        B, T, _ = x.shape
        x = self.input_layer(x)

        # Add CLS token and positional encoding
        cls_token = self.cls_token.repeat(B, 1, 1)
        x = torch.cat([cls_token, x], dim=1)
        x = x + self.pos_embedding[:, : T + 1]

        # Apply Transforrmer
        x = self.dropout(x)
        x = x.transpose(0, 1)
        x = self.transformer(x)

        # Perform classification prediction
        cls = x[0]
        out = self.mlp_head(cls)
        return out


class ViT(L.LightningModule):
    def __init__(self, model_kwargs, lr):
        super().__init__()
        self.save_hyperparameters()
        self.model = VisionTransformer(**model_kwargs)

    def forward(self, x):
        return self.model(x)

    def configure_optimizers(self):
        optimizer = optim.AdamW(self.parameters(), lr=self.hparams.lr)
        lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100, 150], gamma=0.1)
        return [optimizer], [lr_scheduler]

    def _calculate_loss(self, batch, mode="train"):
        imgs, labels = batch
        preds = self.model(imgs)
        loss = F.cross_entropy(preds, labels)
        acc = (preds.argmax(dim=-1) == labels).float().mean()

        self.log("%s_loss" % mode, loss)
        self.log("%s_acc" % mode, acc)
        return loss

    def training_step(self, batch, batch_idx):
        loss = self._calculate_loss(batch, mode="train")
        return loss

    def validation_step(self, batch, batch_idx):
        self._calculate_loss(batch, mode="val")

    def test_step(self, batch, batch_idx):
        self._calculate_loss(batch, mode="test")

在其他文件中引入model.py,实现网络搭建

from model import ViT

model = ViT(model_kwargs={
        "embed_dim": 256,
        "hidden_dim": 512,
        "num_heads": 8,
        "num_layers": 6,
        "patch_size": 4,
        "num_channels": 3,
        "num_patches": 64,
        "num_classes": 10,
        "dropout": 0.2,
        },
        lr=3e-4,
    )

也可以下载预训练的模型

# Files to download
base_url = "https://raw.githubusercontent.com/phlippe/saved_models/main/"
CHECKPOINT_PATH = os.environ.get("PATH_CHECKPOINT", "saved_models/VisionTransformers/")
pretrained_files = [
    "tutorial15/ViT.ckpt",
    "tutorial15/tensorboards/ViT/events.out.tfevents.ViT",
    "tutorial5/tensorboards/ResNet/events.out.tfevents.resnet",
]
# Create checkpoint path if it doesn't exist yet
os.makedirs(CHECKPOINT_PATH, exist_ok=True)

# For each file, check whether it already exists. If not, try downloading it.
for file_name in pretrained_files:
    file_path = os.path.join(CHECKPOINT_PATH, file_name.split("/", 1)[1])
    if "/" in file_name.split("/", 1)[1]:
        os.makedirs(file_path.rsplit("/", 1)[0], exist_ok=True)
    if not os.path.isfile(file_path):
        file_url = base_url + file_name
        print("Downloading %s..." % file_url)
        try:
            urllib.request.urlretrieve(file_url, file_path)
        except HTTPError as e:
            print(
                "Something went wrong. Please try to download the file from the GDrive folder, or contact the author with the full output including the following error:\n",
                e,
            )

pretrained_filename = os.path.join(CHECKPOINT_PATH, "ViT.ckpt")
model = ViT.load_from_checkpoint(pretrained_filename)

torchview.draw_graph 网络可视化

model_graph = draw_graph(model, input_size=(1, 3, 16, 16))
model_graph.resize_graph(scale=5.0)
model_graph.visual_graph.render(format='svg')

运行这段代码,会生成一个svg格式的图片,显示网络结构和每一层的输入输出

训练数据准备

新建一个prepare_data.py

import os
import json
import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms


class CustomDataset(Dataset):
    def __init__(self, image_dir, names, labels, transform=None):
        self.image_dir = image_dir
        self.names = names
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        name_ = self.names[idx]
        img_name = os.path.join(self.image_dir, name_)
        image = Image.open(img_name)

        if self.transform:
            image = self.transform(image)

        label = self.labels[idx]

        return image, label


def load_img_ann(ann_path):
    """return [{img_name, [ (x, y, h, w, label), ... ]}]"""
    with open(ann_path) as fp:
        root = json.load(fp)
    img_dict = {}
    img_label_dict = {}
    for img_info in root['images']:
        img_id = img_info['id']
        img_name = img_info['file_name']
        img_dict[img_id] = {'name': img_name}
    for ann_info in root['annotations']:
        img_id = ann_info['image_id']
        img_category_id = ann_info['category_id']
        img_name = img_dict[img_id]['name']

        img_label_dict[img_id] = {'name': img_name, 'category_id': img_category_id}

    return img_label_dict


def get_dataloader():
    annota_dir = '/home/robotics/Downloads/coco_dataset/annotations/instances_val2017.json'
    img_dir = "/home/robotics/Downloads/coco_dataset/val2017"
    img_dict = load_img_ann(annota_dir)
    values = list(img_dict.values())
    img_names = []
    labels = []
    for item in values:
        category_id = item['category_id']
        labels.append(category_id)
        img_name = item['name']
        img_names.append(img_name)

    # 检查剔除黑白的图片
    img_names_rgb = []
    labels_rgb = []
    for i in range(len(img_names)):
        # 检查文件扩展名,确保它是图片文件(可以根据需要扩展支持的文件类型)
        file_path = os.path.join(img_dir, img_names[i])

        # 打开图片文件
        img = Image.open(file_path)

        # 获取通道数
        num_channels = img.mode
        if num_channels == "RGB" and labels[i] < 10:
            img_names_rgb.append(img_names[i])
            labels_rgb.append(labels[i])

    # 定义一系列图像转换操作
    transform = transforms.Compose([
        transforms.Resize((16, 16)),  # 调整图像大小
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
        transforms.ToTensor(),  # 将图像转换为张量
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化图像
    ])

    # 假设 image_dir 是包含所有图像文件的文件夹路径,labels 是标签列表
    train_set = CustomDataset(img_dir, img_names_rgb[-500:], labels_rgb[-500:], transform=transform)
    val_set = CustomDataset(img_dir, img_names_rgb[-500:-100], labels_rgb[-500:-100], transform=transform)
    test_set = CustomDataset(img_dir, img_names_rgb[-100:], labels_rgb[-100:], transform=transform)

    # 创建一个 DataLoader
    train_loader = DataLoader(train_set, batch_size=32, shuffle=True, drop_last=False)
    val_loader = DataLoader(val_set, batch_size=32, shuffle=True, drop_last=False, num_workers=4)
    test_loader = DataLoader(test_set, batch_size=32, shuffle=True, drop_last=False, num_workers=4)

    return train_loader, val_loader, test_loader


if __name__ == "__main__":
    train_loader, val_loader, test_loader = get_dataloader()
    for batch in train_loader:
        print(batch)

解释一下上面的代码:

这里使用的是coco数据集的2017,可以在官网自行下载,下载下来以后,annotations包含如下内容

这里我们使用的是 instances_val2017.json,如果是正经做训练,应该用train2017,但是train2017文件太大了,处理起来速度很慢,本文仅为说明,不追求训练效果,所以使用val2017进行说明,instances就是用于图像识别的annotation,里面包括了每张图片的label和box,本文创建的ViT 不输出box,仅输出类别。函数

def load_img_ann(ann_path):

是为了将图片的id(每张图片的唯一主键),name和category_id(属于哪一个类别,也就是label)关联起来。

        # 获取通道数
        num_channels = img.mode
        if num_channels == "RGB" and labels[i] < 10:
            img_names_rgb.append(img_names[i])
            labels_rgb.append(labels[i])

注意coco数据集有单通道的黑白图片,要剔除,因为本文的ViT比较简单,输出只能10个类别,所以预处理图片的时候,只选择10个类别。

定义操作变换

    # 定义一系列图像转换操作
    transform = transforms.Compose([
        transforms.Resize((16, 16)),  # 调整图像大小
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
        transforms.ToTensor(),  # 将图像转换为张量
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化图像
    ])

创建一个自己的Dataset类,继承自 torch.utils.data.Dataset

class CustomDataset(Dataset):
    def __init__(self, image_dir, names, labels, transform=None):
        self.image_dir = image_dir
        self.names = names
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        name_ = self.names[idx]
        img_name = os.path.join(self.image_dir, name_)
        image = Image.open(img_name)

        if self.transform:
            image = self.transform(image)

        label = self.labels[idx]

        return image, label

先创建Dataset,再创建dataloader,从Dataset取minibatch。

    # 假设 image_dir 是包含所有图像文件的文件夹路径,labels 是标签列表
    train_set = CustomDataset(img_dir, img_names_rgb[-500:], labels_rgb[-500:], transform=transform)
    val_set = CustomDataset(img_dir, img_names_rgb[-500:-100], labels_rgb[-500:-100], transform=transform)
    test_set = CustomDataset(img_dir, img_names_rgb[-100:], labels_rgb[-100:], transform=transform)

    # 创建一个 DataLoader
    train_loader = DataLoader(train_set, batch_size=32, shuffle=True, drop_last=False)
    val_loader = DataLoader(val_set, batch_size=32, shuffle=True, drop_last=False, num_workers=4)
    test_loader = DataLoader(test_set, batch_size=32, shuffle=True, drop_last=False, num_workers=4)

以上,数据准备工作完成,对模型进行训练

    trainer = L.Trainer(
            default_root_dir=os.path.join(CHECKPOINT_PATH, "ViT"),
            accelerator="auto",
            devices=1,
            max_epochs=180,
            callbacks=[
                ModelCheckpoint(save_weights_only=True, mode="max", monitor="val_acc"),
                LearningRateMonitor("epoch"),
            ],
        )
    trainer.logger._log_graph = True  # If True, we plot the computation graph in tensorboard
    trainer.logger._default_hp_metric = None  # Optional logging argument that we don't need

    trainer.fit(model, train_loader, val_loader)
    # Load best checkpoint after training
    model = ViT.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)

    # Test best model on validation and test set
    val_result = trainer.test(model, dataloaders=val_loader, verbose=False)
    test_result = trainer.test(model, dataloaders=test_loader, verbose=False)
    result = {"test": test_result[0]["test_acc"], "val": val_result[0]["test_acc"]}

完整代码:

一共包括三个文件:model.py 搭建网络的功能, prepare_data.py 数据预处理工作, main.py 网络训练

model.py内容:

import torch.nn as nn
import torch
import torch.optim as optim
import torch.nn.functional as F
import lightning as L


class AttentionBlock(nn.Module):
    def __init__(self, embed_dim, hidden_dim, num_heads, dropout=0.0):
        """
        Inputs:
            embed_dim - Dimensionality of input and attention feature vectors
            hidden_dim - Dimensionality of hidden layer in feed-forward network
                         (usually 2-4x larger than embed_dim)
            num_heads - Number of heads to use in the Multi-Head Attention block
            dropout - Amount of dropout to apply in the feed-forward network
        """
        super().__init__()

        self.layer_norm_1 = nn.LayerNorm(embed_dim)
        self.attn = nn.MultiheadAttention(embed_dim, num_heads)
        self.layer_norm_2 = nn.LayerNorm(embed_dim)
        self.linear = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, embed_dim),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        inp_x = self.layer_norm_1(x)
        x = x + self.attn(inp_x, inp_x, inp_x)[0]
        x = x + self.linear(self.layer_norm_2(x))
        return x


class VisionTransformer(nn.Module):
    def __init__(
        self,
        embed_dim,
        hidden_dim,
        num_channels,
        num_heads,
        num_layers,
        num_classes,
        patch_size,
        num_patches,
        dropout=0.0,
    ):
        """
        Inputs:
            embed_dim - Dimensionality of the input feature vectors to the Transformer
            hidden_dim - Dimensionality of the hidden layer in the feed-forward networks
                         within the Transformer
            num_channels - Number of channels of the input (3 for RGB)
            num_heads - Number of heads to use in the Multi-Head Attention block
            num_layers - Number of layers to use in the Transformer
            num_classes - Number of classes to predict
            patch_size - Number of pixels that the patches have per dimension
            num_patches - Maximum number of patches an image can have
            dropout - Amount of dropout to apply in the feed-forward network and
                      on the input encoding
        """
        super().__init__()

        self.patch_size = patch_size

        # Layers/Networks
        self.input_layer = nn.Linear(num_channels * (patch_size**2), embed_dim)
        self.transformer = nn.Sequential(
            *(AttentionBlock(embed_dim, hidden_dim, num_heads, dropout=dropout) for _ in range(num_layers))
        )
        self.mlp_head = nn.Sequential(nn.LayerNorm(embed_dim), nn.Linear(embed_dim, num_classes))
        self.dropout = nn.Dropout(dropout)

        # Parameters/Embeddings
        self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))
        self.pos_embedding = nn.Parameter(torch.randn(1, 1 + num_patches, embed_dim))

    def img_to_patch(self, x, patch_size, flatten_channels=True):
        """
        Inputs:
            x - Tensor representing the image of shape [B, C, H, W]
            patch_size - Number of pixels per dimension of the patches (integer)
            flatten_channels - If True, the patches will be returned in a flattened format
                               as a feature vector instead of a image grid.
        """
        B, C, H, W = x.shape
        x = x.reshape(B, C, H // patch_size, patch_size, W // patch_size, patch_size)
        x = x.permute(0, 2, 4, 1, 3, 5)  # [B, H', W', C, p_H, p_W]
        x = x.flatten(1, 2)  # [B, H'*W', C, p_H, p_W]
        if flatten_channels:
            x = x.flatten(2, 4)  # [B, H'*W', C*p_H*p_W]
        return x

    def forward(self, x):
        # Preprocess input
        x = self.img_to_patch(x, self.patch_size)
        B, T, _ = x.shape
        x = self.input_layer(x)

        # Add CLS token and positional encoding
        cls_token = self.cls_token.repeat(B, 1, 1)
        x = torch.cat([cls_token, x], dim=1)
        x = x + self.pos_embedding[:, : T + 1]

        # Apply Transforrmer
        x = self.dropout(x)
        x = x.transpose(0, 1)
        x = self.transformer(x)

        # Perform classification prediction
        cls = x[0]
        out = self.mlp_head(cls)
        return out


class ViT(L.LightningModule):
    def __init__(self, model_kwargs, lr):
        super().__init__()
        self.save_hyperparameters()
        self.model = VisionTransformer(**model_kwargs)

    def forward(self, x):
        return self.model(x)

    def configure_optimizers(self):
        optimizer = optim.AdamW(self.parameters(), lr=self.hparams.lr)
        lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100, 150], gamma=0.1)
        return [optimizer], [lr_scheduler]

    def _calculate_loss(self, batch, mode="train"):
        imgs, labels = batch
        preds = self.model(imgs)
        loss = F.cross_entropy(preds, labels)
        acc = (preds.argmax(dim=-1) == labels).float().mean()

        self.log("%s_loss" % mode, loss)
        self.log("%s_acc" % mode, acc)
        return loss

    def training_step(self, batch, batch_idx):
        loss = self._calculate_loss(batch, mode="train")
        return loss

    def validation_step(self, batch, batch_idx):
        self._calculate_loss(batch, mode="val")

    def test_step(self, batch, batch_idx):
        self._calculate_loss(batch, mode="test")

 prepare_data.py内容:

import os
import json
import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms


class CustomDataset(Dataset):
    def __init__(self, image_dir, names, labels, transform=None):
        self.image_dir = image_dir
        self.names = names
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        name_ = self.names[idx]
        img_name = os.path.join(self.image_dir, name_)
        image = Image.open(img_name)

        if self.transform:
            image = self.transform(image)

        label = self.labels[idx]

        return image, label


def load_img_ann(ann_path):
    """return [{img_name, [ (x, y, h, w, label), ... ]}]"""
    with open(ann_path) as fp:
        root = json.load(fp)
    img_dict = {}
    img_label_dict = {}
    for img_info in root['images']:
        img_id = img_info['id']
        img_name = img_info['file_name']
        img_dict[img_id] = {'name': img_name}
    for ann_info in root['annotations']:
        img_id = ann_info['image_id']
        img_category_id = ann_info['category_id']
        img_name = img_dict[img_id]['name']

        img_label_dict[img_id] = {'name': img_name, 'category_id': img_category_id}

    return img_label_dict


def get_dataloader():
    annota_dir = '/home/robotics/Downloads/coco_dataset/annotations/instances_val2017.json'
    img_dir = "/home/robotics/Downloads/coco_dataset/val2017"
    img_dict = load_img_ann(annota_dir)
    values = list(img_dict.values())
    img_names = []
    labels = []
    for item in values:
        category_id = item['category_id']
        labels.append(category_id)
        img_name = item['name']
        img_names.append(img_name)

    # 检查剔除黑白的图片
    img_names_rgb = []
    labels_rgb = []
    for i in range(len(img_names)):
        # 检查文件扩展名,确保它是图片文件(可以根据需要扩展支持的文件类型)
        file_path = os.path.join(img_dir, img_names[i])

        # 打开图片文件
        img = Image.open(file_path)

        # 获取通道数
        num_channels = img.mode
        if num_channels == "RGB" and labels[i] < 10:
            img_names_rgb.append(img_names[i])
            labels_rgb.append(labels[i])

    # 定义一系列图像转换操作
    transform = transforms.Compose([
        transforms.Resize((16, 16)),  # 调整图像大小
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
        transforms.ToTensor(),  # 将图像转换为张量
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化图像
    ])

    # 假设 image_dir 是包含所有图像文件的文件夹路径,labels 是标签列表
    train_set = CustomDataset(img_dir, img_names_rgb[-500:], labels_rgb[-500:], transform=transform)
    val_set = CustomDataset(img_dir, img_names_rgb[-500:-100], labels_rgb[-500:-100], transform=transform)
    test_set = CustomDataset(img_dir, img_names_rgb[-100:], labels_rgb[-100:], transform=transform)

    # 创建一个 DataLoader
    train_loader = DataLoader(train_set, batch_size=32, shuffle=True, drop_last=False)
    val_loader = DataLoader(val_set, batch_size=32, shuffle=True, drop_last=False, num_workers=4)
    test_loader = DataLoader(test_set, batch_size=32, shuffle=True, drop_last=False, num_workers=4)

    return train_loader, val_loader, test_loader


if __name__ == "__main__":
    train_loader, val_loader, test_loader = get_dataloader()
    for batch in train_loader:
        print(batch)

main.py内容:

import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1' # 下面老是报错 shape 不一致
import urllib.request
from urllib.error import HTTPError
import lightning as L
from model import ViT
from torchview import draw_graph
from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint

from prepare_data import get_dataloader


# 加载模型
# Files to download
base_url = "https://raw.githubusercontent.com/phlippe/saved_models/main/"
CHECKPOINT_PATH = os.environ.get("PATH_CHECKPOINT", "saved_models/VisionTransformers/")
pretrained_files = [
    "tutorial15/ViT.ckpt",
    "tutorial15/tensorboards/ViT/events.out.tfevents.ViT",
    "tutorial5/tensorboards/ResNet/events.out.tfevents.resnet",
]
# Create checkpoint path if it doesn't exist yet
os.makedirs(CHECKPOINT_PATH, exist_ok=True)

# For each file, check whether it already exists. If not, try downloading it.
for file_name in pretrained_files:
    file_path = os.path.join(CHECKPOINT_PATH, file_name.split("/", 1)[1])
    if "/" in file_name.split("/", 1)[1]:
        os.makedirs(file_path.rsplit("/", 1)[0], exist_ok=True)
    if not os.path.isfile(file_path):
        file_url = base_url + file_name
        print("Downloading %s..." % file_url)
        try:
            urllib.request.urlretrieve(file_url, file_path)
        except HTTPError as e:
            print(
                "Something went wrong. Please try to download the file from the GDrive folder, or contact the author with the full output including the following error:\n",
                e,
            )

pretrained_filename = os.path.join(CHECKPOINT_PATH, "ViT.ckpt")
needTrain = False
if not os.path.isfile(pretrained_filename):
    print("Found pretrained model at %s, loading..." % pretrained_filename)
    # Automatically loads the model with the saved hyperparameters
    model = ViT.load_from_checkpoint(pretrained_filename)
else:
    L.seed_everything(42)  # To be reproducable
    model = ViT(model_kwargs={
        "embed_dim": 256,
        "hidden_dim": 512,
        "num_heads": 8,
        "num_layers": 6,
        "patch_size": 4,
        "num_channels": 3,
        "num_patches": 64,
        "num_classes": 10,
        "dropout": 0.2,
        },
        lr=3e-4,
    )
    needTrain = True

# 网络结构可视化
model_graph = draw_graph(model, input_size=(1, 3, 16, 16))
model_graph.resize_graph(scale=5.0)
model_graph.visual_graph.render(format='svg')

# 准备训练数据
train_loader, val_loader, test_loader = get_dataloader()

if needTrain:
    trainer = L.Trainer(
            default_root_dir=os.path.join(CHECKPOINT_PATH, "ViT"),
            accelerator="auto",
            devices=1,
            max_epochs=180,
            callbacks=[
                ModelCheckpoint(save_weights_only=True, mode="max", monitor="val_acc"),
                LearningRateMonitor("epoch"),
            ],
        )
    trainer.logger._log_graph = True  # If True, we plot the computation graph in tensorboard
    trainer.logger._default_hp_metric = None  # Optional logging argument that we don't need

    trainer.fit(model, train_loader, val_loader)
    # Load best checkpoint after training
    model = ViT.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)

    # Test best model on validation and test set
    val_result = trainer.test(model, dataloaders=val_loader, verbose=False)
    test_result = trainer.test(model, dataloaders=test_loader, verbose=False)
    result = {"test": test_result[0]["test_acc"], "val": val_result[0]["test_acc"]}

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.coloradmin.cn/o/1164642.html

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈,一经查实,立即删除!

相关文章

【Code Reading】Transformer in vision and video

文章目录 1. vit2. Swin-t3. vit_3D4. TimeSformer First&#x1f680;&#x1f680;5. vivit 1. vit 详细解释 在论文的Table1中有给出三个模型&#xff08;Base/ Large/ Huge&#xff09;的参数&#xff0c;在源码中除了有Patch Size为16x16的外还有32x32的。其中的Layers就…

基于MIMO通信系统的球形译码算法matlab性能仿真,对比PSK检测,SDR检测

目录 1.算法运行效果图预览 2.算法运行软件版本 3.部分核心程序 4.算法理论概述 5.算法完整程序工程 1.算法运行效果图预览 2.算法运行软件版本 MATLAB2022A 3.部分核心程序 ................................................................ for i1:length(SNR) Bit…

SpringCloud(六) Nacos配置管理

Nacos除了可以做注册中心,同样可以做配置管理来使用; Nacos做注册中心的使用和注意事项详看:SpringCloud(四) Nacos注册中心-CSDN博客 目录 一, 统一配置管理 1. 在Nacos中添加配置文件 2. 从微服务拉取配置 (1) 引入nacos-config依赖 (2) 添加bootstrap.yaml (3) 读取n…

相机存储卡被格式化了怎么恢复?数据恢复办法分享!

随着时代的发展&#xff0c;相机被越来越多的用户所使用&#xff0c;这也意味着更多的用户面临着相机数据丢失的问题&#xff0c;很多用户在使用相机的过程中&#xff0c;都出现过不小心格式化相机存储卡的情况&#xff0c;里面的数据也将一并消失&#xff0c;相机存储卡被格式…

Android13充电动画实现

充电动画 以MTK平台为例&#xff0c;实现充电动画 效果图 修改文件清单 system/vendor/mediatek/proprietary/packages/apps/SystemUI/src/com/android/systemui/charging/BubbleBean.javasystem/vendor/mediatek/proprietary/packages/apps/SystemUI/src/com/android/system…

2023年【道路运输企业安全生产管理人员】考试资料及道路运输企业安全生产管理人员考试技巧

题库来源&#xff1a;安全生产模拟考试一点通公众号小程序 道路运输企业安全生产管理人员考试资料根据新道路运输企业安全生产管理人员考试大纲要求&#xff0c;安全生产模拟考试一点通将道路运输企业安全生产管理人员模拟考试试题进行汇编&#xff0c;组成一套道路运输企业安…

浅谈新能源汽车充电桩的选型与安装

叶根胜 安科瑞电气股份有限公司 上海嘉定201801 摘要&#xff1a;电动汽车的大力发展和推广是国家为应对日益突出的燃油供需矛盾和环境污染&#xff0c;加强生态环境保护和治理而开发新能源和清洁能源的措施之一&#xff0c;加快了电动汽车的发展。如今&#xff0c;电动汽车已…

Pure-Pursuit 跟踪双移线 Gazebo 仿真

Pure-Pursuit 跟踪双移线 Gazebo 仿真 主要参考学习下面的博客和开源项目 自动驾驶规划控制&#xff08;&#xff21;*、pure pursuit、LQR算法&#xff0c;使用c在ubuntu和ros环境下实现&#xff09; https://github.com/NeXTzhao/planning Pure-Pursuit 的理论基础见今年六月…

如何在 Endless OS 上安装 ONLYOFFICE 桌面编辑器 7.5

ONLYOFFICE 桌面编辑器是一款基于依据 AGPL v.3 许可进行分发的开源办公套件。使用这款应用&#xff0c;您无需保持网络连接状态即可处理存储在计算机上的文档。本指南会向您介绍&#xff0c;如何在 Endless OS 上安装 ONLYOFFICE 桌面编辑器 7.5。 ONLYOFFICE 桌面版是什么 O…

Ansible中的任务执行控制

循环 简单循环 {{item}} 迭代变量名称 loop: - value1 - value2 - ... //赋值列表{{item}} //迭代变量名称循环散列或字典列表 - name: create filehosts: host1tasks:- name: file moudleservice:name: "{{ item.name }}"state: "{{…

实验记录之——git push

平时做开发的时候经常push代码不成功&#xff0c;如下图 经好友传授经验&#xff0c;有如下方法 Win cmd使用Clash&#xff08;端口是7890&#xff09;代理操作&#xff0c;在cmd中输入&#xff1a; set http_proxy127.0.0.1:7890 set https_proxy127.0.0.1:7890Linux export …

防火墙日志记录和分析

防火墙监控进出网络的流量&#xff0c;并保护部署防火墙的网络免受恶意流量的侵害。它是一个网络安全系统&#xff0c;它根据一些预定义的规则监控传入和传出的流量&#xff0c;它以日志的形式记录有关如何管理流量的信息&#xff0c;日志数据包含流量的源和目标 IP 地址、端口…

可视化流程编排(Bpmn.js)介绍及实践

作者:罗强 为将内部系统打通并规范流程定义&#xff0c;基于统一的平台实现工单自动化流转,从而使用无界山工单系统帮助公司内部人员统一管理和处理来自企业内部提交的工单需求。而在系统中流程编排及节点设计主要是使用bpmn.js实现精细化配置。从而满足各种复杂的业务需求。目…

MapStruct的用法

提示&#xff1a;文章写完后&#xff0c;目录可以自动生成&#xff0c;如何生成可参考右边的帮助文档 文章目录 前言一、依赖导入二、简单使用2.1、定义转换的接口2.2、创建实体类2.3、测试2.4、底层实现 三、常用的注解3.1、Mapping(target "xxx1",source "x…

你还在用System.currentTimeMillis()打印代码执行时间?

文章目录 前言一、开发环境二、使用步骤1. 创建Springboot项目2. 引入hutool3. 使用TimeInterval 总结 前言 Hutool是一个小而全的Java工具类库&#xff0c;里面集成了很多实用的工具类&#xff0c;比如文件、流、加密解密、转码、正则、线程、XML等&#xff0c;通过这些工具类…

探索 Java 8 中的 Stream 流:构建流的多种方式

人嘛&#xff0c;要懂得避嫌… 开篇引入 Java 8引入了Stream流作为一项新的特性&#xff0c;它是用来处理集合数据的一种函数式编程方式。Stream流提供了一种更简洁、高效和易于理解的方法来操作集合数据&#xff0c;同时也能够实现并行处理&#xff0c;以提高性能。 以下是St…

在 GORM 中定义模型

为实现与数据库的无缝交互而打造有效模型的全面指南 在使用 GORM 进行数据库管理时&#xff0c;定义模型是基础。模型是您的应用程序的面向对象结构与数据库的关系世界之间的桥梁。本文深入探讨了在 GORM 中打造有效模型的艺术&#xff0c;探讨如何创建结构化的 Go 结构体&…

第十二章,集合类例题

例题1 package 例题;import java.util.*;public class 例题 {public static void main(String[] args) {// TODO Auto-generated method stub//实例化集合类对象Collection<String> list new ArrayList<>();//调用方法&#xff0c;向集合添加数据list.add("…

Java——java.time包使用方法详解

Java——time包使用方法详解 java.time 包是 Java 8 引入的新日期和时间 API&#xff08;JSR 310&#xff09;&#xff0c;用于替代旧的 java.util.Date 和 java.util.Calendar 类。它提供了一组全新的类来处理日期、时间、时间间隔、时区等&#xff0c;具有更好的设计和易用性…

什么是消息队列

什么是消息队列 消息队列是一种通信机制&#xff0c;用于在不同的应用程序或组件之间传递消息。它允许应用程序之间异步地发送和接收消息&#xff0c;而无需直接依赖彼此的可用性或性能。消息队列通常用于解耦不同组件&#xff0c;提高系统的可伸缩性和可维护性&#xff0c;以…