常见ML任务的模型和模型实用程序
xt-models的Python项目详细描述
xt型号
说明
这个repo包含处理ML任务的通用模型和实用程序,由Xtract AI开发。在
还有更多。在
安装
来自PyPi:
pip install xt-models
来源:
^{pr2}$使用
获取分割模型
fromxt_models.modelsimportModelBuilder,SegmentationModulefromtorchimportnndeep_sup_scale=0.4fc_dim=2048n_class=2net_encoder=ModelBuilder.build_encoder(arch="resnet50dilated",fc_dim=fc_dim,weights="/nasty/scratch/common/smart_objects/model/ade20k/encoder_epoch_20.pth")net_decoder=ModelBuilder.build_decoder(arch="ppm_deepsup",fc_dim=fc_dim,num_class=150,weights="/nasty/scratch/common/smart_objects/model/ade20k/decoder_epoch_20.pth")in_channels=net_decoder.conv_last[-1].in_channelsnet_decoder.conv_last[-1]=nn.Conv2d(in_channels,n_class,kernel_size=(1,1),stride=(1,1))net_decoder.conv_last_deepsup=nn.Conv2d(in_channels,n_class,1,1,0)model=SegmentationModule(net_encoder,net_decoder,deep_sup_scale)
获取检测模型
from xt_models.models import Model
import torch
# Load a fine-tuned model for inference
model_name = "yolov5x"
model = Model(model_name,nc=15)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
weights = "/nasty/scratch/common/smart_objects/model/veh_detection/yolov5_ft/best_state_dict.pt"
ckpt = torch.load(weights, map_location=device)
model.load_state_dict(ckpt['model_state_dict'])
# Load pre-trained COCO model for finetuning/inference
model_name = "yolov5x"
model = Model(model_name,nc=80)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
weights = "/nasty/scratch/common/smart_objects/model/veh_detection/yolov5_pretrain/yolov5x_state_dict.pt"
ckpt = torch.load(weights, map_location=device)
model.load_state_dict(ckpt['model_state_dict'])
# Fine-tuning number of classes
n_class = 15
model.nc = n_class
实施新模型
如果您总是要为不同的项目复制和粘贴相同的模型代码,只需将模型代码添加到models
目录中,并将其导入models/__init__.py
文件。在
数据源
[数据说明和链接]
依赖关系/许可
[依赖项及其许可证的列表,包括数据]
参考文献
[参考文献列表]
- 项目
标签: