0
点赞
收藏
分享

微信扫一扫

对yolov5未训练完剔除节点/层

洲行 2022-01-18 阅读 97
import glob
import math
import os
import random
import shutil
import subprocess
import time
from copy import copy
from pathlib import Path
from sys import platform

import cv2
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torchvision
import yaml
from scipy.signal import butter, filtfilt
from tqdm import tqdm

from . import torch_utils  #  torch_utils, google_utils




def create_pretrained(f='weights/best.pt', s='weights/pretrained.pt'):  # from utils.utils import *; create_pretrained()
    # create pretrained checkpoint 's' from 'f' (create_pretrained(x, x) for x in glob.glob('./*.pt'))
    x = torch.load(f, map_location=torch.device('cpu'))
    x['optimizer'] = None
    x['training_results'] = None
    x['epoch'] = -1
    x['model'].half()  # to FP16
    x['ema'] = None
    for p in x['model'].parameters():
        p.requires_grad = True
    torch.save(x, s)
    print('%s saved as pretrained checkpoint %s, %.1fMB' % (f, s, os.path.getsize(s) / 1E6))

以后代码执行完毕后,就可以进行迁移学习训练。多类迁移到少类或者少类迁移到多类。

官方util.general里面的对于最终训练完的进行无关节点剔除:

删除anchor_grid:

举报

相关推荐

0 条评论