TensorFlow Hub (2)-Object Detection Colab

TensorFlow Hub (2)-Object Detection Colab

导入和设置

import os
import pathlib

import matplotlib
import matplotlib.pyplot as plt

import io
import scipy.misc
import numpy as np
from six import BytesIO
from PIL import Image, ImageDraw, ImageFont
from six.moves.urllib.request import urlopen

import tensorflow as tf
import tensorflow_hub as hub

tf.get_logger().setLevel('ERROR')

实用工具函数

运行以下单元格以创建稍后将需要的一些实用程序:
包括load_image_into_numpy_array 将图片转换为numpy array
模型的字典
IMAGES_FOR_TEST 图片名称和地址的字典
COCO17_HUMAN_POSE_KEYPOINTS 人体姿态关键点 的list

加载图像的辅助方法utilities

Map of Model Name to TF Hub handle 模型名称到TF Hub handle的映射
COCO 2017数据集的带有人类关键点的元组列表。具有关键点的模型需要此功能。

# @title Run this!!

def load_image_into_numpy_array(path):
  """Load an image from file into a numpy array.

  Puts image into numpy array to feed into tensorflow graph.
  Note that by convention we put it into a numpy array with shape
  (height, width, channels), where channels=3 for RGB.

  Args:
    path: the file path to the image

  Returns:
    uint8 numpy array with shape (img_height, img_width, 3)
  """
  image = None
  if(path.startswith('http')):
    response = urlopen(path)
    image_data = response.read()
    image_data = BytesIO(image_data)
    image = Image.open(image_data)
  else:
    image_data = tf.io.gfile.GFile(path, 'rb').read()
    image = Image.open(BytesIO(image_data))

  (im_width, im_height) = image.size
  return np.array(image.getdata()).reshape(
      (1, im_height, im_width, 3)).astype(np.uint8)


ALL_MODELS = {
   
'CenterNet HourGlass104 512x512' : 'https://hub.tensorflow.google.cn/tensorflow/centernet/hourglass_512x512/1',
'CenterNet HourGlass104 Keypoints 512x512' : 'https://hub.tensorflow.google.cn/tensorflow/centernet/hourglass_512x512_kpts/1',
'CenterNet HourGlass104 1024x1024' : 'https://hub.tensorflow.google.cn/tensorflow/centernet/hourglass_1024x1024/1',
'CenterNet HourGlass104 Keypoints 1024x1024' : 'https://hub.tensorflow.google.cn/tensorflow/centernet/hourglass_1024x1024_kpts/1',
'CenterNet Resnet50 V1 FPN 512x512' : 'https://hub.tensorflow.google.cn/tensorflow/centernet/resnet50v1_fpn_512x512/1',
'CenterNet Resnet50 V1 FPN Keypoints 512x512' : 'https://hub.tensorflow.google.cn/tensorflow/centernet/resnet50v1_fpn_512x512_kpts/1',
'CenterNet Resnet101 V1 FPN 512x512' : 'https://hub.tensorflow.google.cn/tensorflow/centernet/resnet101v1_fpn_512x512/1',
'CenterNet Resnet50 V2 512x512' : 'https://hub.tensorflow.google.cn/tensorflow/centernet/resnet50v2_512x512/1',
'CenterNet Resnet50 V2 Keypoints 512x512' : 'https://hub.tensorflow.google.cn/tensorflow/centernet/resnet50v2_512x512_kpts/1',
'EfficientDet D0 512x512' : 'https://hub.tensorflow.google.cn/tensorflow/efficientdet/d0/1',
'EfficientDet D1 640x640' : 'https://hub.tensorflow.google.cn/tensorflow/efficientdet/d1/1',
'EfficientDet D2 768x768' : 'https://hub.tensorflow.google.cn/tensorflow/efficientdet/d2/1',
'EfficientDet D3 896x896' : 'https://hub.tensorflow.google.cn/tensorflow/efficientdet/d3/1',
'EfficientDet D4 1024x1024' : 'https://hub.tensorflow.google.cn/tensorflow/efficientdet/d4/1',
'EfficientDet D5 1280x1280' : 'https://hub.tensorflow.google.cn/tensorflow/efficientdet/d5/1',
'EfficientDet D6 1280x1280' : 'https://hub.tensorflow.google.cn/tensorflow/efficientdet/d6/1',
'EfficientDet D7 1536x1536' : 'https://hub.tensorflow.google.cn/tensorflow/efficientdet/d7/1',
'
  • 2
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 3
    评论
Sure! Here are the steps to fine-tune ViT-S on a custom dataset using Google Colab: 1. Open a new Google Colab notebook and select a GPU runtime environment. 2. Install the necessary libraries: ``` !pip install torch torchvision !pip install timm ``` 3. Download and prepare the custom dataset. You can use any dataset of your choice. Make sure to split it into training and validation sets. 4. Define the data loaders: ``` import torch import torchvision.transforms as transforms from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder # Define the transformations transform_train = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) transform_val = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # Define the data loaders train_dataset = ImageFolder('path_to_train_data', transform=transform_train) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4) val_dataset = ImageFolder('path_to_val_data', transform=transform_val) val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4) ``` Replace 'path_to_train_data' and 'path_to_val_data' with the paths to your training and validation data folders, respectively. 5. Load the pre-trained ViT-S model: ``` import timm model = timm.create_model('vit_small_patch16_224', pretrained=True) ``` 6. Modify the last layer of the model to fit your custom dataset: ``` import torch.nn as nn num_classes = len(train_dataset.classes) model.head = nn.Sequential( nn.LayerNorm((768,)), nn.Linear(768, num_classes) ) ``` Replace '768' with the hidden size of the model you are using. For ViT-S, it is 768. 7. Define the optimizer and criterion: ``` import torch.optim as optim optimizer = optim.Adam(model.parameters(), lr=1e-4) criterion = nn.CrossEntropyLoss() ``` 8. Fine-tune the model: ``` device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) num_epochs = 10 for epoch in range(num_epochs): train_loss = 0.0 val_loss = 0.0 correct = 0 total = 0 # Train the model model.train() for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() train_loss += loss.item() * inputs.size(0) # Evaluate the model on validation set model.eval() with torch.no_grad(): for inputs, labels in val_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) loss = criterion(outputs, labels) val_loss += loss.item() * inputs.size(0) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() train_loss = train_loss / len(train_loader.dataset) val_loss = val_loss / len(val_loader.dataset) accuracy = 100 * correct / total print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f} \tAccuracy: {:.2f}'.format( epoch+1, train_loss, val_loss, accuracy)) ``` 9. Save the model: ``` torch.save(model.state_dict(), 'path_to_save_model') ``` Replace 'path_to_save_model' with the path where you want to save the model. That's it! You have fine-tuned ViT-S on your custom dataset using Google Colab.
评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值