初始化换发型项目：3个微服务代码 + 部署脚本
包含： - hair_service_sd: 换发型/换发色算法服务 (端口 8801) - photo_service: LoRA 训练调度服务 (端口 32678) - stable-diffusion-webui: SD WebUI 推理服务 (端口 57860) - kohya_ss_home: 训练环境代码 - meidaojia: 监控测试脚本 - setup.sh: 一键部署脚本 (conda环境恢复 + 配置生成 + 完整性检查) - start_all_services.sh: 启动3个服务 - configure.ini.template: 路径模板化 (BASE_DIR自动推导) - conda_envs/py310.yml: py310 环境定义 大文件 (weights/, models/, data/, conda_envs/*.tar.gz 等) 通过 .gitignore 排除， 由网盘单独上传。
This commit is contained in:
colomi
@@ -0,0 +1,222 @@
|
||||
import os
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from seg.networks.vgg import vgg16
|
||||
|
||||
__all__ = ['get_fcn32s', 'get_fcn16s', 'get_fcn8s',
|
||||
'get_fcn32s_vgg16_voc', 'get_fcn16s_vgg16_voc', 'get_fcn8s_vgg16_voc']
|
||||
|
||||
|
||||
class FCN32s(nn.Module):
|
||||
"""There are some difference from original fcn"""
|
||||
|
||||
def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True,
|
||||
norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(FCN32s, self).__init__()
|
||||
self.aux = aux
|
||||
if backbone == 'vgg16':
|
||||
self.pretrained = vgg16(pretrained=pretrained_base).features
|
||||
else:
|
||||
raise RuntimeError('unknown backbone: {}'.format(backbone))
|
||||
self.head = _FCNHead(512, nclass, norm_layer)
|
||||
if aux:
|
||||
self.auxlayer = _FCNHead(512, nclass, norm_layer)
|
||||
|
||||
self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head'])
|
||||
|
||||
def forward(self, x):
|
||||
size = x.size()[2:]
|
||||
pool5 = self.pretrained(x)
|
||||
|
||||
outputs = []
|
||||
out = self.head(pool5)
|
||||
out = F.interpolate(out, size, mode='bilinear', align_corners=True)
|
||||
outputs.append(out)
|
||||
|
||||
if self.aux:
|
||||
auxout = self.auxlayer(pool5)
|
||||
auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True)
|
||||
outputs.append(auxout)
|
||||
|
||||
return tuple(outputs)
|
||||
|
||||
|
||||
class FCN16s(nn.Module):
|
||||
def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(FCN16s, self).__init__()
|
||||
self.aux = aux
|
||||
if backbone == 'vgg16':
|
||||
self.pretrained = vgg16(pretrained=pretrained_base).features
|
||||
else:
|
||||
raise RuntimeError('unknown backbone: {}'.format(backbone))
|
||||
self.pool4 = nn.Sequential(*self.pretrained[:24])
|
||||
self.pool5 = nn.Sequential(*self.pretrained[24:])
|
||||
self.head = _FCNHead(512, nclass, norm_layer)
|
||||
self.score_pool4 = nn.Conv2d(512, nclass, 1)
|
||||
if aux:
|
||||
self.auxlayer = _FCNHead(512, nclass, norm_layer)
|
||||
|
||||
self.__setattr__('exclusive', ['head', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool4'])
|
||||
|
||||
def forward(self, x):
|
||||
pool4 = self.pool4(x)
|
||||
pool5 = self.pool5(pool4)
|
||||
|
||||
outputs = []
|
||||
score_fr = self.head(pool5)
|
||||
|
||||
score_pool4 = self.score_pool4(pool4)
|
||||
|
||||
upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True)
|
||||
fuse_pool4 = upscore2 + score_pool4
|
||||
|
||||
out = F.interpolate(fuse_pool4, x.size()[2:], mode='bilinear', align_corners=True)
|
||||
outputs.append(out)
|
||||
|
||||
if self.aux:
|
||||
auxout = self.auxlayer(pool5)
|
||||
auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True)
|
||||
outputs.append(auxout)
|
||||
|
||||
return tuple(outputs)
|
||||
|
||||
|
||||
class FCN8s(nn.Module):
|
||||
def __init__(self, nclass, backbone='vgg16', aux=False, pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(FCN8s, self).__init__()
|
||||
self.aux = aux
|
||||
if backbone == 'vgg16':
|
||||
self.pretrained = vgg16(pretrained=pretrained_base).features
|
||||
else:
|
||||
raise RuntimeError('unknown backbone: {}'.format(backbone))
|
||||
self.pool3 = nn.Sequential(*self.pretrained[:17])
|
||||
self.pool4 = nn.Sequential(*self.pretrained[17:24])
|
||||
self.pool5 = nn.Sequential(*self.pretrained[24:])
|
||||
self.head = _FCNHead(512, nclass, norm_layer)
|
||||
self.score_pool3 = nn.Conv2d(256, nclass, 1)
|
||||
self.score_pool4 = nn.Conv2d(512, nclass, 1)
|
||||
if aux:
|
||||
self.auxlayer = _FCNHead(512, nclass, norm_layer)
|
||||
|
||||
self.__setattr__('exclusive',
|
||||
['head', 'score_pool3', 'score_pool4', 'auxlayer'] if aux else ['head', 'score_pool3',
|
||||
'score_pool4'])
|
||||
|
||||
def forward(self, x):
|
||||
pool3 = self.pool3(x)
|
||||
pool4 = self.pool4(pool3)
|
||||
pool5 = self.pool5(pool4)
|
||||
|
||||
outputs = []
|
||||
score_fr = self.head(pool5)
|
||||
|
||||
score_pool4 = self.score_pool4(pool4)
|
||||
score_pool3 = self.score_pool3(pool3)
|
||||
|
||||
upscore2 = F.interpolate(score_fr, score_pool4.size()[2:], mode='bilinear', align_corners=True)
|
||||
fuse_pool4 = upscore2 + score_pool4
|
||||
|
||||
upscore_pool4 = F.interpolate(fuse_pool4, score_pool3.size()[2:], mode='bilinear', align_corners=True)
|
||||
fuse_pool3 = upscore_pool4 + score_pool3
|
||||
|
||||
out = F.interpolate(fuse_pool3, x.size()[2:], mode='bilinear', align_corners=True)
|
||||
outputs.append(out)
|
||||
|
||||
if self.aux:
|
||||
auxout = self.auxlayer(pool5)
|
||||
auxout = F.interpolate(auxout, x.size()[2:], mode='bilinear', align_corners=True)
|
||||
outputs.append(auxout)
|
||||
|
||||
return tuple(outputs)
|
||||
|
||||
|
||||
class _FCNHead(nn.Module):
|
||||
def __init__(self, in_channels, channels, norm_layer=nn.BatchNorm2d, **kwargs):
|
||||
super(_FCNHead, self).__init__()
|
||||
inter_channels = in_channels // 4
|
||||
self.block = nn.Sequential(
|
||||
nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
|
||||
norm_layer(inter_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout(0.1),
|
||||
nn.Conv2d(inter_channels, channels, 1)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.block(x)
|
||||
|
||||
|
||||
def get_fcn32s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models',
|
||||
pretrained_base=True, **kwargs):
|
||||
acronyms = {
|
||||
'pascal_voc': 'pascal_voc',
|
||||
'pascal_aug': 'pascal_aug',
|
||||
'ade20k': 'ade',
|
||||
'coco': 'coco',
|
||||
'citys': 'citys',
|
||||
}
|
||||
from ..data.dataloader import datasets
|
||||
model = FCN32s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
|
||||
if pretrained:
|
||||
from .model_store import get_model_file
|
||||
device = torch.device(kwargs['local_rank'])
|
||||
model.load_state_dict(torch.load(get_model_file('fcn32s_%s_%s' % (backbone, acronyms[dataset]), root=root),
|
||||
map_location=device))
|
||||
return model
|
||||
|
||||
|
||||
def get_fcn16s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models',
|
||||
pretrained_base=True, **kwargs):
|
||||
acronyms = {
|
||||
'pascal_voc': 'pascal_voc',
|
||||
'pascal_aug': 'pascal_aug',
|
||||
'ade20k': 'ade',
|
||||
'coco': 'coco',
|
||||
'citys': 'citys',
|
||||
}
|
||||
from ..data.dataloader import datasets
|
||||
model = FCN16s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
|
||||
if pretrained:
|
||||
from .model_store import get_model_file
|
||||
device = torch.device(kwargs['local_rank'])
|
||||
model.load_state_dict(torch.load(get_model_file('fcn16s_%s_%s' % (backbone, acronyms[dataset]), root=root),
|
||||
map_location=device))
|
||||
return model
|
||||
|
||||
|
||||
def get_fcn8s(dataset='pascal_voc', backbone='vgg16', pretrained=False, root='~/.torch/models',
|
||||
pretrained_base=True, **kwargs):
|
||||
acronyms = {
|
||||
'pascal_voc': 'pascal_voc',
|
||||
'pascal_aug': 'pascal_aug',
|
||||
'ade20k': 'ade',
|
||||
'coco': 'coco',
|
||||
'citys': 'citys',
|
||||
}
|
||||
from ..data.dataloader import datasets
|
||||
model = FCN8s(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
|
||||
if pretrained:
|
||||
from .model_store import get_model_file
|
||||
device = torch.device(kwargs['local_rank'])
|
||||
model.load_state_dict(torch.load(get_model_file('fcn8s_%s_%s' % (backbone, acronyms[dataset]), root=root),
|
||||
map_location=device))
|
||||
return model
|
||||
|
||||
|
||||
def get_fcn32s_vgg16_voc(**kwargs):
|
||||
return get_fcn32s('pascal_voc', 'vgg16', **kwargs)
|
||||
|
||||
|
||||
def get_fcn16s_vgg16_voc(**kwargs):
|
||||
return get_fcn16s('pascal_voc', 'vgg16', **kwargs)
|
||||
|
||||
|
||||
def get_fcn8s_vgg16_voc(**kwargs):
|
||||
return get_fcn8s('pascal_voc', 'vgg16', **kwargs)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
model = FCN16s(21)
|
||||
print(model)
|
||||
Reference in New Issue
Block a user