聚类项目：2.模型搭建，开始训练_sovit 聚类模型

【选中CPU/GPU】

# 选择设备（CPU/GPU），这里用yolov5源码，健壮性强
def select_device(device=''):
    # device = None or 'cpu' or 0 or '0' or '1'
    device = str(device).strip().lower().replace('cuda:', '').replace('none', '')  # to string, 'cuda:0' to '0'
    cpu = device == 'cpu'
    if cpu:
        os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
    elif device:
        os.environ['CUDA_VISIBLE_DEVICES'] = device
        assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \
            f"Invalid CUDA '--device {device}'requested, use'--device cpu' or pass valid CUDA device(s)"
    if not cpu and torch.cuda.is_available():
        devices = device.split(',') if device else '0'
        n = len(devices)  # 显卡块数
        assert n == 1, 'only support 1 gpu accelaration'
        arg = 'cuda:0'
    else:
        arg = 'cpu'

    return torch.device(arg)


device = select_device(device)
gpu_flag = device.type == 'cuda'
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

【模型设计：直接调用vgg16抽取特征】

# feat model . torchvision为我们封装好了常见的模型。具体方法和变量用法见下面
# 1.调用torchvision.models.vgg16():https://blog.csdn.net/u014380165/article/details/79119664
# 2.1 model.eval()不启用bn和dropout,用在测试中。  https://www.jianshu.com/p/778dc3b96141
# 2.2 model.train()则启用bn和dropout,用在训练中。 https://blog.csdn.net/qq_38410428/article/details/101102075
# 3.vgg16().features: 输出模型全部结构，每一层固定序号  https://blog.csdn.net/agoodboy1997/article/details/110959360
# 4.to(device)：把模型送入cuda
feat_model = tv.models.vgg16(weights='DEFAULT').eval().features.to(device)
1
2
3
4
5
6

【对转换成mp4格式的视频进行抽帧和切片处理】