def __init__(self, config, device, **kwargs):
super(Net_small, self).__init__()
layers_ = config.MODEL.LAYERS
pretrained = config.MODEL.PRETRAINED
classes = config.MODEL.classes
atrous_rates = config.MODEL.atrous_rates
self.device = device
if layers_ == 50:
resnet = model.resnet50(pretrained=pretrained)
elif layers_ == 34:
resnet = model.resnet34(pretrained=pretrained)
else:
resnet = model.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1_my, resnet.bn1, resnet.relu)
self.max_pool = resnet.maxpool
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
# del resnet
self.relu = nn.ReLU()
self.confidence = Confidence(classes, 1)
self.aspp = ASPP(in_channels=512, atrous_rates=atrous_rates)
self.cls = nn.Sequential(
nn.Conv2d(256, 256, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(256), nn.ReLU(inplace=True),
nn.Conv2d(256, classes, kernel_size=1))
initialize_weights(self.confidence)
initialize_weights(self.cls)
initialize_weights(self.aspp)
def predict():
# fasle 返回 未训练的模型
predict_model = resnet152_bo(resnet152(pretrained=True))
predict_model.load_state_dict(
torch.load(opt.load_model_path,
map_location=lambda storage, loc: storage))
# 模型改为预测模式
predict_model.eval()
# 如果GPU可用，加载到GPU
if opt.use_gpu:
predict_model.cuda()
# 测试集照片地址
test_img_dir = opt.test_img_dir
image = cv2.imread(test_img_dir)
# result中内容为 左上角坐标、右下角坐标、类别名、输入的测试图地址、预测类别的可能性
result = predict_result(predict_model, test_img_dir)
for left_up, right_bottom, class_name, _, prob in result:
# 将预测框添加到测试图片中
cv2.rectangle(image, left_up, right_bottom, (0, 255, 0), 2)
# 预测框的左上角写入 所属类别
cv2.putText(image, class_name, left_up, cv2.FONT_HERSHEY_SIMPLEX, 1,
(0, 0, 255), 1, cv2.LINE_AA)
print(prob)
# 将测试结果写入
cv2.imwrite(opt.result_img_dir, image)
def model_selection(self):
if self.config.model_name == "resnet18":
return resnet18(in_channels=int(self.input_shape[0]),
num_classes=self.label_num)
elif self.config.model_name == "resnet34":
return resnet34(in_channels=int(self.input_shape[0]),
num_classes=self.label_num)
elif self.config.model_name == "resnet50":
return resnet50(in_channels=int(self.input_shape[0]),
num_classes=self.label_num)
elif self.config.model_name == "resnet101":
return resnet101(in_channels=int(self.input_shape[0]),
num_classes=self.label_num)
elif self.config.model_name == "resnet152":
return resnet152(in_channels=int(self.input_shape[0]),
num_classes=self.label_num)
elif self.config.model_name == "convnet":
return convnet(in_channels=int(self.input_shape[0]),
num_classes=self.label_num)
elif self.config.model_name == "resnet50":
return wide_resnet50_2(in_channels=int(self.input_shape[0]),
num_classes=self.label_num)
elif self.config.model_name == "resnet101":
return wide_resnet101_2(in_channels=int(self.input_shape[0]),
num_classes=self.label_num)
def get_network(args, model_args, use_gpu=True):
""" return given network
"""
if args.net == 'resnet18':
from models.resnet import resnet18
net = resnet18(**model_args)
elif args.net == 'resnet34':
from models.resnet import resnet34
net = resnet34(**model_args)
elif args.net == 'resnet50':
from models.resnet import resnet50
net = resnet50(**model_args)
elif args.net == 'resnet101':
from models.resnet import resnet101
net = resnet101(**model_args)
elif args.net == 'resnet152':
from models.resnet import resnet152
net = resnet152(**model_args)
else:
print('the network name you have entered is not supported yet')
sys.exit()
if use_gpu:
net = net.cuda()
return net
def get_model(train_model):
if train_model == 'resnet18':
return resnet.resnet18()
elif train_model == 'resnet34':
return resnet.resnet34()
elif train_model == 'resnet50':
return resnet.resnet50()
elif train_model == 'resnet101':
return resnet.resnet101()
elif train_model == 'resnet152':
return resnet.resnet152()
elif train_model == 'resnet18_copy':
return resnet_copy.resnet18()
elif train_model == 'resnet34_copy':
return resnet_copy.resnet34()
elif train_model == 'resnet50_copy':
return resnet_copy.resnet50()
elif train_model == 'resnet101_copy':
return resnet_copy.resnet101()
elif train_model == 'resnet152':
return resnet_copy.resnet152()
elif train_model == 'vgg11':
return vgg11()
elif train_model == 'vgg13':
return vgg13()
elif train_model == 'vgg16':
return vgg16()
elif train_model == 'vgg19':
return vgg19()
elif train_model == 'nin':
return nin()
elif train_model == 'googlenet':
return googlenet()
def Resnet(opt):
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
if opt.model_depth == 10:
model = resnet.resnet10(
num_classes=opt.n_classes)
elif opt.model_depth == 18:
model = resnet.resnet18(
num_classes=opt.n_classes,
pool=opt.pool)
elif opt.model_depth == 34:
model = resnet.resnet34(
num_classes=opt.n_classes,
pool=opt.pool)
elif opt.model_depth == 50:
model = resnet.resnet50(
num_classes=opt.n_classes,
pool=opt.pool)
elif opt.model_depth == 101:
model = resnet.resnet101(
num_classes=opt.n_classes)
elif opt.model_depth == 152:
model = resnet.resnet152(
num_classes=opt.n_classes)
elif opt.model_depth == 200:
model = resnet.resnet200(
num_classes=opt.n_classes)
return model
def get_edge_scores(junctions, regions, rgb_folder, _id, epoch=1, model='resnet152'):
# load model
resnet = resnet152(pretrained=False).cuda()
edge_classifier = EdgeClassifier(resnet)
edge_classifier = edge_classifier.cuda()
edge_classifier = edge_classifier.eval()
# open RGB image
split = 'det'
out_size = 256
rgb_path = os.path.join(rgb_folder, _id +'.jpg')
rgb = Image.open(rgb_path).resize((out_size, out_size))
rgb = np.array(rgb)/255.0
model_path = '/home/nelson/Workspace/building_reconstruction/working_model/binary_edge_classifier_with_regions/saved_models/edge_classifier_{}_{}_iter_{}.pth'.format(model, split, epoch)
edge_classifier.load_state_dict(torch.load(model_path))
# check backup -- save time
temp_dir = './temp/{}/'.format(model)
if not os.path.isdir(temp_dir):
os.makedirs(temp_dir)
backup_path = '{}/{}_{}_{}.npy'.format(temp_dir, _id, epoch, split)
if os.path.isfile(backup_path):
return np.load(open(backup_path, 'rb'), encoding='bytes').item()
# combine regions
all_reg = np.zeros((out_size, out_size))
for k, reg in enumerate(regions):
reg = Image.fromarray(reg*255.0).resize((out_size, out_size))
reg = np.array(reg)/255.0
inds = np.array(np.where(reg==1))
all_reg[inds[0, :], inds[1, :]] = k
# draw edge
lw_from_cls = {}
for k, c1 in enumerate(junctions):
for l, c2 in enumerate(junctions):
if k > l:
edge = Image.new('L', (out_size, out_size))
dr = ImageDraw.Draw(edge)
x1, y1 = c1
x2, y2 = c2
div = 256.0/out_size
dr.line((int(x1/div), int(y1/div), int(x2/div), int(y2/div)), fill="white", width=int(4/div))
edge = np.array(edge)/255.0
imgs = np.concatenate([rgb.transpose(2, 0, 1), edge[np.newaxis, :, :], all_reg[np.newaxis, :, :]], 0)
imgs = torch.from_numpy(imgs).cuda().float()
prob = edge_classifier(imgs.unsqueeze(0))
prob = prob.detach().cpu().numpy()
lw_from_cls[(k, l)] = prob[0]
# save backup
np.save(open(backup_path, 'wb'), lw_from_cls)
return lw_from_cls
def __init__(self, config, **kwargs):
super(Net, self).__init__()
layers_ = config.MODEL.LAYERS
pretrained = config.MODEL.PRETRAINED
classes = config.MODEL.classes
atrous_rates = config.MODEL.atrous_rates
if layers_ == 50:
resnet = model.resnet50(pretrained=pretrained)
elif layers_ == 101:
resnet = model.resnet101(pretrained=pretrained)
else:
resnet = model.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.conv2, resnet.bn2, resnet.relu,
resnet.conv3, resnet.bn3, resnet.relu)
self.max_pool = resnet.maxpool
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
# del resnet
self.side_output0 = SideOutput(128, num_class=classes)
self.side_output1 = SideOutput(256, num_class=classes)
self.side_output2 = SideOutput(512, num_class=classes)
# self.side_output3 = SideOutput(1024, 1, kernel_sz=15, stride=8, padding=7)
self.side_output4 = SideOutput(2048, num_class=classes)
self.sigmoid = nn.Sigmoid()
self.fuse = nn.Conv2d(4, 1, kernel_size=1, bias=False)
self.relu = nn.ReLU()
self.confidence = Confidence(classes, 1)
self.cw = nn.Conv2d(2, 1, kernel_size=1, padding=0, bias=False)
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
self.aspp = ASPP(in_channels=2048 + 256, atrous_rates=atrous_rates)
self.bot_fine = nn.Conv2d(256, 48, kernel_size=1, bias=False)
self.cls = nn.Sequential(
nn.Conv2d(256 + 48, 256, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(256), nn.ReLU(inplace=True),
nn.Conv2d(256, classes, kernel_size=1))
initialize_weights(self.confidence)
initialize_weights(self.cls)
def initialize_model(model_name, num_dense_layers=2, dropout=0):
'''
Initialise a model with a custom head
to predict both sequence length and digits
Parameters
----------
model_name : str
Model Name can be either:
ResNet
VGG
ConvNet
BaselineCNN_dropout
Returns
-------
model : object
The model to be initialize
'''
if model_name[:3] == "VGG":
model = VGG(model_name, num_classes=7)
model.classifier = CustomHead(512)
elif model_name[:6] == "ResNet":
if model_name == "ResNet18":
model = resnet18(num_classes=7)
model.linear = CustomHead(512)
elif model_name == "ResNet34":
model = resnet18(num_classes=7)
model.linear = CustomHead(512)
elif model_name == "ResNet50":
model = resnet50(num_classes=7)
model.linear = CustomHead(512 * 4)
elif model_name == "ResNet101":
model = resnet101(num_classes=7)
model.linear = CustomHead(512 * 4)
elif model_name == "ResNet152":
model = resnet152(num_classes=7)
model.linear = CustomHead(512 * 4)
elif model_name == "ConvNet":
model = ConvModel(num_dense_layers=num_dense_layers, dropout=dropout)
elif model_name == "BaselineCNNdropout":
model = BaselineCNNdropout(num_classes=7, p=dropout)
model.fc2 = CustomHead(4096)
return model
def generate_model(opt):
assert opt.model in ['resnet', 'densenet', 'se_resnet']
if opt.model == 'resnet':
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
from models.resnet import get_fine_tuning_parameters
if opt.model_depth == 10:
model = resnet.resnet10(pretrained=True, num_classes=opt.n_classes)
elif opt.model_depth == 18:
model = resnet.resnet18(pretrained=True, num_classes=opt.n_classes)
elif opt.model_depth == 34:
model = resnet.resnet34(pretrained=True, num_classes=opt.n_classes)
elif opt.model_depth == 50:
model = resnet.resnet50(pretrained=True, num_classes=opt.n_classes)
elif opt.model_depth == 101:
model = resnet.resnet101(pretrained=True,
num_classes=opt.n_classes)
elif opt.model_depth == 152:
model = resnet.resnet152(pretrained=True,
num_classes=opt.n_classes)
elif opt.model_depth == 200:
model = resnet.resnet200(pretrained=True,
num_classes=opt.n_classes)
elif opt.model == 'se_resnet':
assert opt.model_depth in [18, 34, 50, 101, 152]
from models.se_resnet import get_fine_tuning_parameters
if opt.model_depth == 18:
model = se_resnet.resnet18(pretrained=True,
num_classes=opt.n_classes)
elif opt.model_depth == 34:
model = se_resnet.resnet34(pretrained=True,
num_classes=opt.n_classes)
elif opt.model_depth == 50:
model = se_resnet.resnet50(pretrained=True,
num_classes=opt.n_classes)
elif opt.model_depth == 101:
model = se_resnet.resnet101(pretrained=True,
num_classes=opt.n_classes)
if not opt.no_cuda:
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
return model, parameters
else:
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
return model, parameters
return model, model.parameters()
def get_model_param(args):
# assert args.model in ['resnet', 'vgg']
if args.model == 'resnet':
assert args.model_depth in [18, 34, 50, 101, 152]
from models.resnet import get_fine_tuning_parameters
if args.model_depth == 18:
model = resnet.resnet18(pretrained=False,
input_size=args.input_size,
num_classes=args.n_classes)
elif args.model_depth == 34:
model = resnet.resnet34(pretrained=False,
input_size=args.input_size,
num_classes=args.n_classes)
elif args.model_depth == 50:
model = resnet.resnet50(pretrained=False,
input_size=args.input_size,
num_classes=args.n_classes)
elif args.model_depth == 101:
model = resnet.resnet101(pretrained=False,
input_size=args.input_size,
num_classes=args.n_classes)
elif args.model_depth == 152:
model = resnet.resnet152(pretrained=False,
input_size=args.input_size,
num_classes=args.n_classes)
# elif args.model == 'vgg':
# pass
# Load pretrained model here
if args.finetune:
pretrained_model = model_path[args.arch]
args.pretrain_path = os.path.join(args.root_path, 'pretrained_models',
pretrained_model)
print("=> loading pretrained model '{}'...".format(pretrained_model))
model.load_state_dict(torch.load(args.pretrain_path))
# Only modify the last layer
if args.model == 'resnet':
model.fc = nn.Linear(model.fc.in_features, args.n_finetune_classes)
# elif args.model == 'vgg':
# pass
parameters = get_fine_tuning_parameters(model, args.ft_begin_index,
args.lr_mult1, args.lr_mult2)
return model, parameters
return model, model.parameters()
def get_network(args):
if args.net == 'vgg16':
from models.vgg import vgg16
model_ft = vgg16(args.num_classes, export_onnx=args.export_onnx)
elif args.net == 'alexnet':
from models.alexnet import alexnet
model_ft = alexnet(num_classes=args.num_classes,
export_onnx=args.export_onnx)
elif args.net == 'mobilenet':
from models.mobilenet import mobilenet_v2
model_ft = mobilenet_v2(pretrained=True, export_onnx=args.export_onnx)
elif args.net == 'vgg19':
from models.vgg import vgg19
model_ft = vgg19(args.num_classes, export_onnx=args.export_onnx)
else:
if args.net == 'googlenet':
from models.googlenet import googlenet
model_ft = googlenet(pretrained=True)
elif args.net == 'inception':
from models.inception import inception_v3
model_ft = inception_v3(args,
pretrained=True,
export_onnx=args.export_onnx)
elif args.net == 'resnet18':
from models.resnet import resnet18
model_ft = resnet18(pretrained=True, export_onnx=args.export_onnx)
elif args.net == 'resnet34':
from models.resnet import resnet34
model_ft = resnet34(pretrained=True, export_onnx=args.export_onnx)
elif args.net == 'resnet101':
from models.resnet import resnet101
model_ft = resnet101(pretrained=True, export_onnx=args.export_onnx)
elif args.net == 'resnet50':
from models.resnet import resnet50
model_ft = resnet50(pretrained=True, export_onnx=args.export_onnx)
elif args.net == 'resnet152':
from models.resnet import resnet152
model_ft = resnet152(pretrained=True, export_onnx=args.export_onnx)
else:
print("The %s is not supported..." % (args.net))
return
if args.net == 'mobilenet':
num_ftrs = model_ft.classifier[1].in_features
model_ft.classifier[1] = nn.Linear(num_ftrs * 4, args.num_classes)
else:
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, args.num_classes)
net = model_ft
return net
def __init__(self,
layers=50,
bins=(1, 2, 3, 6),
dropout=0.1,
classes=21,
use_ppm=True,
pretrained=True):
super(PSPNet, self).__init__()
assert layers in [50, 101, 152]
assert 2048 % len(bins) == 0
assert classes > 1
self.use_ppm = use_ppm
if layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.conv2, resnet.bn2, resnet.relu,
resnet.conv3, resnet.bn3, resnet.relu,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
fea_dim = 2048
if use_ppm:
self.ppm = PPM(fea_dim, int(fea_dim / len(bins)), bins)
fea_dim *= 2
self.cls = nn.Sequential(
nn.Conv2d(fea_dim, 512, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(512), nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout), nn.Conv2d(512, classes, kernel_size=1))
if self.training:
self.aux = nn.Sequential(
nn.Conv2d(1024, 256, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(256), nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout), nn.Conv2d(256, classes,
kernel_size=1))
def get_network(args,cfg):
""" return given network
"""
# pdb.set_trace()
if args.net == 'lenet5':
net = LeNet5().cuda()
elif args.net == 'alexnet':
net = alexnet(pretrained=args.pretrain, num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg16':
net = vgg16(pretrained=args.pretrain, num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg13':
net = vgg13(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg11':
net = vgg11(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg19':
net = vgg19(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg16_bn':
net = vgg16_bn(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg13_bn':
net = vgg13_bn(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg11_bn':
net = vgg11_bn(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'vgg19_bn':
net = vgg19_bn(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net =='inceptionv3':
net = inception_v3().cuda()
# elif args.net == 'inceptionv4':
# net = inceptionv4().cuda()
# elif args.net == 'inceptionresnetv2':
# net = inception_resnet_v2().cuda()
elif args.net == 'resnet18':
net = resnet18(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda(args.gpuid)
elif args.net == 'resnet34':
net = resnet34(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'resnet50':
net = resnet50(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda(args.gpuid)
elif args.net == 'resnet101':
net = resnet101(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'resnet152':
net = resnet152(pretrained=args.pretrain,num_classes=cfg.PARA.train.num_classes).cuda()
elif args.net == 'squeezenet':
net = squeezenet1_0().cuda()
else:
print('the network name you have entered is not supported yet')
sys.exit()
return net
def _get_model(self, backbone):
if backbone == 'resnet18':
model = resnet18(pretrained=True, num_classes=self.args.classnum).to(self.args.device)
elif backbone == 'resnet34':
model = resnet34(pretrained=True, num_classes=self.args.classnum).to(self.args.device)
elif backbone == 'resnet50':
model = resnet50(pretrained=True, num_classes=self.args.classnum).to(self.args.device)
elif backbone == 'resnet101':
model = resnet101(pretrained=True, num_classes=self.args.classnum).to(self.args.device)
elif backbone == 'resnet152':
model = resnet152(pretrained=True, num_classes=self.args.classnum).to(self.args.device)
elif backbone == 'preact_resnet18':
model = PreActResNet18(num_classes=self.args.classnum, input_size=self.args.image_size,
input_dim=self.args.input_dim).to(self.args.device)
elif backbone == 'preact_resnet34':
model = PreActResNet34(num_classes=self.args.classnum, input_size=self.args.image_size,
input_dim=self.args.input_dim).to(self.args.device)
elif backbone == 'preact_resnet50':
model = PreActResNet50(num_classes=self.args.classnum, input_size=self.args.image_size,
input_dim=self.args.input_dim).to(self.args.device)
elif backbone == 'preact_resnet101':
model = PreActResNet101(num_classes=self.args.classnum, input_size=self.args.image_size,
input_dim=self.args.input_dim).to(self.args.device)
elif backbone == 'preact_resnet152':
model = PreActResNet152(num_classes=self.args.classnum, input_size=self.args.image_size,
input_dim=self.args.input_dim).to(self.args.device)
elif backbone == 'densenet121':
model = densenet121(num_classes=self.args.classnum, pretrained=True).to(self.args.device)
elif backbone == 'densenet161':
model = densenet161(num_classes=self.args.classnum, pretrained=True).to(self.args.device)
elif backbone == 'densenet169':
model = densenet169(num_classes=self.args.classnum, pretrained=True).to(self.args.device)
elif backbone == 'densenet201':
model = densenet201(num_classes=self.args.classnum, pretrained=True).to(self.args.device)
elif backbone == 'mlp':
model = MLPNet().to(self.args.device)
elif backbone == 'cnn_small' or backbone == "CNN_SMALL":
model = CNN_small(self.args.classnum).to(self.args.device)
elif backbone == "cnn" or backbone == "CNN":
model = CNN(n_outputs=self.args.classnum, input_channel=self.args.input_dim, linear_num=self.args.linear_num).to(self.args.device)
else:
print("No matched backbone. Using ResNet50...")
model = resnet50(pretrained=True, num_classes=self.args.classnum,
input_size=self.args.image_size).to(self.args.device)
return model
def create_landmarks_model(landmarks_model_path,
landmarks_network='resnet_50',
landmarks_num_classes=196,
landmarks_img_size=256):
use_cuda = torch.cuda.is_available()
#---------------------------------------------------------------- 构建 landmarks 模型
if landmarks_network == 'resnet_18':
landmarks_model = resnet18(num_classes=landmarks_num_classes,
img_size=landmarks_img_size)
elif landmarks_network == 'resnet_34':
landmarks_model = resnet34(num_classes=landmarks_num_classes,
img_size=landmarks_img_size)
elif landmarks_network == 'resnet_50':
landmarks_model = resnet50(num_classes=landmarks_num_classes,
img_size=landmarks_img_size)
elif landmarks_network == 'resnet_101':
landmarks_model = resnet101(num_classes=landmarks_num_classes,
img_size=landmarks_img_size)
elif landmarks_network == 'resnet_152':
landmarks_model = resnet152(num_classes=landmarks_num_classes,
img_size=landmarks_img_size)
elif landmarks_network == 'mobilenetv2':
landmarks_model = MobileNetV2(n_class=ops.landmarks_num_classes,
input_size=ops.landmarks_img_size[0])
else:
print('error no the struct model : {}'.format(ops.model))
device = torch.device("cuda:0" if use_cuda else "cpu")
# 加载测试模型
if os.access(landmarks_model_path, os.F_OK): # checkpoint
# chkpt = torch.load(ops.landmarks_model, map_location=device)
# landmarks_model.load_state_dict(chkpt)
chkpt = torch.load(landmarks_model_path,
map_location=lambda storage, loc: storage)
landmarks_model.load_state_dict(chkpt)
landmarks_model.eval() # 设置为前向推断模式
print('load landmarks model : {}'.format(landmarks_model_path))
print(
'\n/******************* landmarks model acc ******************/')
acc_model('', landmarks_model)
landmarks_model = landmarks_model.to(device)
return landmarks_model
def get_network(args):
""" return given network
"""
if args.task == 'cifar10':
nclass = 10
elif args.task == 'cifar100':
nclass = 100
#Yang added none bn vggs
if args.net == 'vgg11':
from models.vgg import vgg11
net = vgg11(num_classes=nclass)
elif args.net == 'vgg13':
from models.vgg import vgg13
net = vgg13(num_classes=nclass)
elif args.net == 'vgg16':
from models.vgg import vgg16
net = vgg16(num_classes=nclass)
elif args.net == 'vgg19':
from models.vgg import vgg19
net = vgg19(num_classes=nclass)
elif args.net == 'resnet18':
from models.resnet import resnet18
net = resnet18(num_classes=nclass)
elif args.net == 'resnet34':
from models.resnet import resnet34
net = resnet34(num_classes=nclass)
elif args.net == 'resnet50':
from models.resnet import resnet50
net = resnet50(num_classes=nclass)
elif args.net == 'resnet101':
from models.resnet import resnet101
net = resnet101(num_classes=nclass)
elif args.net == 'resnet152':
from models.resnet import resnet152
net = resnet152(num_classes=nclass)
else:
print('the network name you have entered is not supported yet')
sys.exit()
if args.gpu: #use_gpu
net = net.cuda()
return net
def get_resnet_3d(num_classes=2,
model_depth=10,
shortcut_type='B',
sample_size=112,
sample_duration=16):
assert model_depth in [10, 18, 34, 50, 101, 152, 200]
if model_depth == 10:
model = resnet.resnet10(num_classes=num_classes,
shortcut_type=shortcut_type,
sample_size=sample_size,
sample_duration=sample_duration)
elif model_depth == 18:
model = resnet.resnet18(num_classes=num_classes,
shortcut_type=shortcut_type,
sample_size=sample_size,
sample_duration=sample_duration)
elif model_depth == 34:
model = resnet.resnet34(num_classes=num_classes,
shortcut_type=shortcut_type,
sample_size=sample_size,
sample_duration=sample_duration)
elif model_depth == 50:
model = resnet.resnet50(num_classes=num_classes,
shortcut_type=shortcut_type,
sample_size=sample_size,
sample_duration=sample_duration)
elif model_depth == 101:
model = resnet.resnet101(num_classes=num_classes,
shortcut_type=shortcut_type,
sample_size=sample_size,
sample_duration=sample_duration)
elif model_depth == 152:
model = resnet.resnet152(num_classes=num_classes,
shortcut_type=shortcut_type,
sample_size=sample_size,
sample_duration=sample_duration)
else:
model = resnet.resnet200(num_classes=num_classes,
shortcut_type=shortcut_type,
sample_size=sample_size,
sample_duration=sample_duration)
return model
def get_model(self):
self.model_fname = os.listdir(self.config.load_path + "model/")[0]
self.model_type = self.model_fname.split("_")[0] # ex: resnet50
if self.model_type == "resnet18":
return resnet18(int(self.input_shape[0]), self.label_num)
elif self.model_type == "resnet34":
return resnet34(int(self.input_shape[0]), self.label_num)
elif self.model_type == "resnet50":
return resnet50(int(self.input_shape[0]), self.label_num)
elif self.model_type == "resnet101":
return resnet101(int(self.input_shape[0]), self.label_num)
elif self.model_type == "resnet152":
return resnet152(int(self.input_shape[0]), self.label_num)
elif self.model_type == "resnet50":
return wide_resnet50_2(in_channels=int(self.input_shape[0]),
num_classes=self.label_num)
elif self.model_type == "resnet101":
return wide_resnet101_2(in_channels=int(self.input_shape[0]),
num_classes=self.label_num)
def get_network(args):
# return given network
if args.net == 'resnet18':
from models.resnet import resnet18
net = resnet18()
elif args.net == 'resnet34':
from models.resnet import resnet34
net = resnet34()
elif args.net == 'resnet50':
from models.resnet import resnet50
net = resnet50()
elif args.net == 'resnet101':
from models.resnet import resnet101
net = resnet101()
elif args.net == 'resnet152':
from models.resnet import resnet152
net = resnet152()
else:
print('the network name you have entered is not supported yet')
sys.exit()
return net
def generate_model(opt):
assert opt.model in ['resnet', 'resnext']
if opt.model == 'resnet':
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
from models.resnet import get_fine_tuning_parameters
if opt.model_depth == 10:
model = resnet.resnet10(opt=opt)
elif opt.model_depth == 18:
model = resnet.resnet18(opt=opt)
elif opt.model_depth == 34:
model = resnet.resnet34(opt=opt)
elif opt.model_depth == 50:
model = resnet.resnet50(opt=opt)
elif opt.model_depth == 101:
model = resnet.resnet101(opt=opt)
elif opt.model_depth == 152:
model = resnet.resnet152(opt=opt)
elif opt.model_depth == 200:
model = resnet.resnet200(opt=opt)
elif opt.model == 'resnext':
assert opt.model_depth in [50, 101, 152]
from models.resnext import get_fine_tuning_parameters
if opt.model_depth == 50:
model = resnext.resnet50(opt=opt)
elif opt.model_depth == 101:
model = resnext.resnet101(opt=opt)
elif opt.model_depth == 152:
model = resnext.resnet152(opt=opt)
if not opt.no_cuda:
model = model.cuda()
return model, model.parameters()
def create_model(model):
print('==> 2. Building the Model...')
if model == 'lenet':
net = lenet(dataset.num_classes, dataset.img_channels).to(device)
elif model == 'convnet':
net = convnet(dataset.num_classes, dataset.img_channels).to(device)
elif model == 'resnet32_wide':
net = resnet32_wide(dataset.num_classes, dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'densenet40':
net = densenet40(num_classes = dataset.num_classes, in_channels = dataset.img_channels).to(device)
elif model == 'densenet161':
net = densenet161(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'resnet50':
net = resnet50(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'resnet110':
net = resnet110(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'resnet152':
net = resnet152(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'resnet110_SD':
net = resnet110_SD(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
elif model == 'resnet152_SD':
net = resnet152_SD(num_classes = dataset.num_classes, in_channels = dataset.img_channels, pretrained = dataset.require_pretrained).to(device)
return net
def generate_model(opt):
assert opt.model in [
'resnet'
]
if opt.model == 'resnet':
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
if opt.model_depth == 10:
model = resnet.resnet10()
elif opt.model_depth == 18:
model = resnet.resnet18()
elif opt.model_depth == 34:
model = resnet.resnet34()
elif opt.model_depth == 50:
model = resnet.resnet50()
elif opt.model_depth == 101:
model = resnet.resnet101()
elif opt.model_depth == 152:
model = resnet.resnet152()
elif opt.model_depth == 200:
model = resnet.resnet200()
return model
def generate_model(opt):
assert opt.model in [
'resnet', 'preresnet', 'wideresnet', 'resnext', 'densenet'
]
if opt.model == 'resnet':
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
from models.resnet import get_fine_tuning_parameters
if opt.model_depth == 10:
model = resnet.resnet10(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 18:
model = resnet.resnet18(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 34:
model = resnet.resnet34(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 50:
model = resnet.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
model_type=opt.model_type)
elif opt.model_depth == 101:
model = resnet.resnet101(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
model_type=opt.model_type)
elif opt.model_depth == 152:
model = resnet.resnet152(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
model_type=opt.model_type)
elif opt.model_depth == 200:
model = resnet.resnet200(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
model_type=opt.model_type)
elif opt.model == 'wideresnet':
assert opt.model_depth in [50]
from models.wide_resnet import get_fine_tuning_parameters
if opt.model_depth == 50:
model = wide_resnet.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
k=opt.wide_resnet_k,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'resnext':
assert opt.model_depth in [50, 101, 152]
from models.resnext import get_fine_tuning_parameters
if opt.model_depth == 50:
model = resnext.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 101:
model = resnext.resnet101(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 152:
model = resnext.resnet152(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'preresnet':
assert opt.model_depth in [18, 34, 50, 101, 152, 200]
from models.pre_act_resnet import get_fine_tuning_parameters
if opt.model_depth == 18:
model = pre_act_resnet.resnet18(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 34:
model = pre_act_resnet.resnet34(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 50:
model = pre_act_resnet.resnet50(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 101:
model = pre_act_resnet.resnet101(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 152:
model = pre_act_resnet.resnet152(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 200:
model = pre_act_resnet.resnet200(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'densenet':
assert opt.model_depth in [121, 169, 201, 264]
from models.densenet import get_fine_tuning_parameters
if opt.model_depth == 121:
model = densenet.densenet121(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 169:
model = densenet.densenet169(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 201:
model = densenet.densenet201(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 264:
model = densenet.densenet264(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
if not opt.no_cuda:
import os
# os.environ['CUDA_VISIBLE_DEVICES'] = f'{opt.cuda_id}'
model = model.cuda(device=opt.cuda_id)
model = nn.DataParallel(model, device_ids=[0]) # CUDA change
if opt.pretrain_path:
print('loading pretrained model {}'.format(opt.pretrain_path))
pretrain = torch.load(opt.pretrain_path)
print(pretrain['arch'])
arch = f'{opt.model}-{opt.model_depth}'
# arch = opt.model + '-' + opt.model_depth
print(arch)
assert arch == pretrain['arch']
model.load_state_dict(pretrain['state_dict'])
if opt.model == 'densenet':
model.module.classifier = nn.Linear(
model.module.classifier.in_features,
opt.n_finetune_classes)
model.module.classifier = model.module.classifier.cuda(
device=opt.cuda_id)
# elif opt.use_quadriplet:
# model = EmbeddingModel(model, opt.n_finetune_classes, not opt.no_cuda, opt.cuda_id)
else:
model.module.fc = nn.Sequential(
nn.Dropout(0.4),
nn.Linear(model.module.fc.in_features, 512), nn.ReLU6(),
nn.Dropout(0.4), nn.Linear(512, 128), nn.ReLU6(),
nn.Linear(128,
opt.n_finetune_classes)).cuda(device=opt.cuda_id)
# model.module.fc = nn.Linear(model.module.fc.in_features,
# opt.n_finetune_classes)
# model.module.fc = model.module.fc.cuda(device=opt.cuda_id)
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
print(len(list(parameters)), 'params to fine tune')
return model, parameters
else:
if opt.pretrain_path:
print('loading pretrained model {}'.format(opt.pretrain_path))
pretrain = torch.load(opt.pretrain_path)
assert opt.arch == pretrain['arch']
model.load_state_dict(pretrain['state_dict'])
if opt.model == 'densenet':
model.classifier = nn.Linear(model.classifier.in_features,
opt.n_finetune_classes)
else:
model.fc = nn.Linear(model.fc.in_features,
opt.n_finetune_classes)
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
return model, parameters
return model, model.parameters()
from models.resnet import resnet152
from iceberg import *
from torch.optim import SGD
lr = 0.0005
mom = 0.9
wd = 1e-4
bsize = 172
model = resnet152(True)
model = torch.nn.DataParallel(model).cuda()
optim = SGD(model.parameters(), lr, momentum=mom, weight_decay=wd)
iceberg = Iceberg('./data/train.json',
'./data/resnet152_lr5e-4_mom9e-1_wd1e-4_bs{}_pretrained_model.pth'.format(bsize), model, optim, 10000, bsize)
iceberg.run(transforms.Compose([
transforms.Resize(256),
transforms.RandomCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()]))
def get_model(class_num):
if (MODEL_TYPE == 'alexnet'):
model = alexnet.alexnet(pretrained=FINETUNE)
elif (MODEL_TYPE == 'vgg'):
if (MODEL_DEPTH_OR_VERSION == 11):
model = vgg.vgg11(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 13):
model = vgg.vgg13(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 16):
model = vgg.vgg16(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 19):
model = vgg.vgg19(pretrained=FINETUNE)
else:
print('Error : VGG should have depth of either [11, 13, 16, 19]')
sys.exit(1)
elif (MODEL_TYPE == 'squeezenet'):
if (MODEL_DEPTH_OR_VERSION == 0 or MODEL_DEPTH_OR_VERSION == 'v0'):
model = squeezenet.squeezenet1_0(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 1 or MODEL_DEPTH_OR_VERSION == 'v1'):
model = squeezenet.squeezenet1_1(pretrained=FINETUNE)
else:
print('Error : Squeezenet should have version of either [0, 1]')
sys.exit(1)
elif (MODEL_TYPE == 'resnet'):
if (MODEL_DEPTH_OR_VERSION == 18):
model = resnet.resnet18(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 34):
model = resnet.resnet34(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 50):
model = resnet.resnet50(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 101):
model = resnet.resnet101(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 152):
model = resnet.resnet152(pretrained=FINETUNE)
else:
print(
'Error : Resnet should have depth of either [18, 34, 50, 101, 152]'
)
sys.exit(1)
elif (MODEL_TYPE == 'densenet'):
if (MODEL_DEPTH_OR_VERSION == 121):
model = densenet.densenet121(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 169):
model = densenet.densenet169(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 161):
model = densenet.densenet161(pretrained=FINETUNE)
elif (MODEL_DEPTH_OR_VERSION == 201):
model = densenet.densenet201(pretrained=FINETUNE)
else:
print(
'Error : Densenet should have depth of either [121, 169, 161, 201]'
)
sys.exit(1)
elif (MODEL_TYPE == 'inception'):
if (MODEL_DEPTH_OR_VERSION == 3 or MODEL_DEPTH_OR_VERSION == 'v3'):
model = inception.inception_v3(pretrained=FINETUNE)
else:
print('Error : Inception should have version of either [3, ]')
sys.exit(1)
else:
print(
'Error : Network should be either [alexnet / squeezenet / vgg / resnet / densenet / inception]'
)
sys.exit(1)
if (MODEL_TYPE == 'alexnet' or MODEL_TYPE == 'vgg'):
num_ftrs = model.classifier[6].in_features
feature_model = list(model.classifier.children())
feature_model.pop()
feature_model.append(nn.Linear(num_ftrs, class_num))
model.classifier = nn.Sequential(*feature_model)
elif (MODEL_TYPE == 'resnet' or MODEL_TYPE == 'inception'):
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, class_num)
elif (MODEL_TYPE == 'densenet'):
num_ftrs = model.classifier.in_features
model.classifier = nn.Linear(num_ftrs, class_num)
return model
def generate_model(opt):
assert opt.model in [
'resnet', 'preresnet', 'wideresnet', 'resnext', 'densenet', 'i3d',
'i3dv2'
]
if opt.model == 'resnet':
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
from models.resnet import get_fine_tuning_parameters
if opt.model_depth == 10:
model = resnet.resnet10(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 18:
model = resnet.resnet18(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 34:
model = resnet.resnet34(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 50:
model = resnet.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 101:
model = resnet.resnet101(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 152:
model = resnet.resnet152(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 200:
model = resnet.resnet200(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'wideresnet':
assert opt.model_depth in [50]
from models.wide_resnet import get_fine_tuning_parameters
if opt.model_depth == 50:
model = wide_resnet.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
k=opt.wide_resnet_k,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'resnext':
assert opt.model_depth in [50, 101, 152]
from models.resnext import get_fine_tuning_parameters
if opt.model_depth == 50:
model = resnext.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 101:
model = resnext.resnet101(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 152:
model = resnext.resnet152(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'preresnet':
assert opt.model_depth in [18, 34, 50, 101, 152, 200]
from models.pre_act_resnet import get_fine_tuning_parameters
if opt.model_depth == 18:
model = pre_act_resnet.resnet18(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 34:
model = pre_act_resnet.resnet34(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 50:
model = pre_act_resnet.resnet50(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 101:
model = pre_act_resnet.resnet101(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 152:
model = pre_act_resnet.resnet152(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 200:
model = pre_act_resnet.resnet200(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'densenet':
assert opt.model_depth in [121, 169, 201, 264]
from models.densenet import get_fine_tuning_parameters
if opt.model_depth == 121:
model = densenet.densenet121(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 169:
model = densenet.densenet169(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 201:
model = densenet.densenet201(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 264:
model = densenet.densenet264(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == "i3d":
from models.i3dpt import get_fine_tuning_parameters
model = i3dpt.I3D(num_classes=opt.n_classes, dropout_prob=0.5)
elif opt.model == "i3dv2":
from models.I3D_Pytorch import get_fine_tuning_parameters
model = I3D_Pytorch.I3D(num_classes=opt.n_classes,
dropout_keep_prob=0.5)
if not opt.no_cuda:
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
if opt.pretrain_path:
print('loading pretrained model {}'.format(opt.pretrain_path))
pretrain = torch.load(opt.pretrain_path)
if opt.model != "i3d" and opt.model != "i3dv2":
assert opt.arch == pretrain['arch']
model.load_state_dict(pretrain['state_dict'])
else:
pretrain = {"module." + k: v for k, v in pretrain.items()}
model_dict = model.state_dict()
model_dict.update(pretrain)
model.load_state_dict(model_dict)
if opt.model == 'densenet':
model.module.classifier = nn.Linear(
model.module.classifier.in_features,
opt.n_finetune_classes)
model.module.classifier = model.module.classifier.cuda()
else:
model.module.fc = nn.Linear(model.module.fc.in_features,
opt.n_finetune_classes)
model.module.fc = model.module.fc.cuda()
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
return model, parameters
else:
if opt.pretrain_path:
print('loading pretrained model {}'.format(opt.pretrain_path))
pretrain = torch.load(opt.pretrain_path)
assert opt.arch == pretrain['arch']
model.load_state_dict(pretrain['state_dict'])
if opt.model == 'densenet':
model.classifier = nn.Linear(model.classifier.in_features,
opt.n_finetune_classes)
else:
model.fc = nn.Linear(model.fc.in_features,
opt.n_finetune_classes)
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
return model, parameters
return model, model.parameters()
def generate_model(opt):
assert opt.mode in ['score', 'feature']
if opt.mode == 'score':
last_fc = True
elif opt.mode == 'feature':
last_fc = False
assert opt.model_name in [
'resnet', 'preresnet', 'wideresnet', 'resnext', 'densenet'
]
if opt.model_name == 'resnet':
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
if opt.model_depth == 10:
model = resnet.resnet10(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 18:
model = resnet.resnet18(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 34:
model = resnet.resnet34(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 50:
model = resnet.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 101:
model = resnet.resnet101(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 152:
model = resnet.resnet152(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 200:
model = resnet.resnet200(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_name == 'wideresnet':
assert opt.model_depth in [50]
if opt.model_depth == 50:
model = wide_resnet.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
k=opt.wide_resnet_k,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_name == 'resnext':
assert opt.model_depth in [50, 101, 152]
if opt.model_depth == 50:
model = resnext.resnet50(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 101:
model = resnext.resnet101(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 152:
model = resnext.resnet152(num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_name == 'preresnet':
assert opt.model_depth in [18, 34, 50, 101, 152, 200]
if opt.model_depth == 18:
model = pre_act_resnet.resnet18(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 34:
model = pre_act_resnet.resnet34(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 50:
model = pre_act_resnet.resnet50(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 101:
model = pre_act_resnet.resnet101(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 152:
model = pre_act_resnet.resnet152(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 200:
model = pre_act_resnet.resnet200(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_name == 'densenet':
assert opt.model_depth in [121, 169, 201, 264]
if opt.model_depth == 121:
model = densenet.densenet121(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 169:
model = densenet.densenet169(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 201:
model = densenet.densenet201(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
elif opt.model_depth == 264:
model = densenet.densenet264(num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
last_fc=last_fc)
if not opt.no_cuda:
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
return model
print('{} train iters per epoch:'.format(len(trainloader)))
############ choose loss and backbone ######################
if opt.loss == 'focal_loss':
criterion = focal_loss.FocalLoss(gamma=2)
else:
criterion = torch.nn.CrossEntropyLoss()
if opt.backbone == 'resnet18':
model = resnet.resnet_face18(use_se=opt.use_se)
elif opt.backbone == 'resnet50':
model = resnet.resnet50()
elif opt.backbone == 'resnet101':
model = resnet.resnet101()
elif opt.backbone == 'resnet152':
model = resnet.resnet152()
elif opt.backbone == 'googlenet':
model = googlenet.GoogLeNet()
if opt.metric == 'add_margin':
metric_fc = metrics.AddMarginProduct(1024,
opt.num_classes,
s=30,
m=0.35)
elif opt.metric == 'arc_margin':
metric_fc = metrics.ArcMarginProduct(512,
opt.num_classes,
s=30,
m=0.5,
easy_margin=opt.easy_margin)
elif opt.metric == 'sphere':
def generate_model(opt):
assert opt.model in [
'resnet', 'resnext'
]
if opt.model == 'resnet':
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
from models.resnet import get_fine_tuning_parameters
if opt.model_depth == 10:
model = resnet.resnet10(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
isSource = opt.isSource,
transfer_module = opt.transfer_module,
sourceKind = opt.sourceKind,
layer_num = opt.layer_num,
multi_source = opt.multi_source)
elif opt.model_depth == 18:
model = resnet.resnet18(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
isSource = opt.isSource,
transfer_module = opt.transfer_module,
sourceKind = opt.sourceKind,
layer_num = opt.layer_num,
multi_source = opt.multi_source)
elif opt.model_depth == 34:
model = resnet.resnet34(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
isSource = opt.isSource,
transfer_module = opt.transfer_module,
sourceKind = opt.sourceKind,
layer_num = opt.layer_num,
multi_source = opt.multi_source)
elif opt.model_depth == 50:
model = resnet.resnet50(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
isSource = opt.isSource,
transfer_module = opt.transfer_module,
sourceKind = opt.sourceKind,
layer_num = opt.layer_num,
multi_source = opt.multi_source)
elif opt.model_depth == 101:
model = resnet.resnet101(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
isSource = opt.isSource,
transfer_module = opt.transfer_module,
sourceKind = opt.sourceKind,
layer_num = opt.layer_num,
multi_source = opt.multi_source)
elif opt.model_depth == 152:
model = resnet.resnet152(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
isSource = opt.isSource,
transfer_module = opt.transfer_module,
sourceKind = opt.sourceKind,
layer_num = opt.layer_num,
multi_source = opt.multi_source)
elif opt.model_depth == 200:
model = resnet.resnet200(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
isSource = opt.isSource,
transfer_module = opt.transfer_module,
sourceKind = opt.sourceKind,
layer_num = opt.layer_num,
multi_source = opt.multi_source)
elif opt.model == 'resnext':
assert opt.model_depth in [50, 101, 152]
from models.resnext import get_fine_tuning_parameters
if opt.model_depth == 50:
model = resnext.resnet50(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
isSource = opt.isSource,
transfer_module = opt.transfer_module,
sourceKind = opt.sourceKind,
layer_num = opt.layer_num,
multi_source = opt.multi_source)
elif opt.model_depth == 101:
model = resnext.resnet101(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
isSource = opt.isSource,
transfer_module = opt.transfer_module,
sourceKind = opt.sourceKind,
layer_num = opt.layer_num,
multi_source = opt.multi_source)
elif opt.model_depth == 152:
model = resnext.resnet152(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration,
isSource = opt.isSource,
transfer_module = opt.transfer_module,
sourceKind = opt.sourceKind,
layer_num = opt.layer_num,
multi_source = opt.multi_source)
print(opt.no_cuda)
print(type(opt.no_cuda))
if not opt.no_cuda:
if opt.pretrain_path:
print('loading pretrained model {}'.format(opt.pretrain_path))
pretrain = torch.load(opt.pretrain_path)
print('loading pretrained model arch', pretrain['arch'], opt.arch)
assert opt.arch == pretrain['arch']
pretrained_dict = pretrain['state_dict']
model_dict = model.state_dict()
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
pretrained_dict = {str.replace(k,'module.',''): v for k,v in pretrained_dict.items()}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
if opt.inference == False:
model.module.fc = nn.Linear(model.module.fc.in_features,
opt.n_finetune_classes)
model.module.fc = model.module.fc.cuda()
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
print(model)
return model, parameters
elif opt.inference:
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
return model, model.parameters()
else:
if opt.pretrain_path:
print('loading pretrained model {}'.format(opt.pretrain_path))
print('loading pretrained model arch', pretrain['arch'])
pretrain = torch.load(opt.pretrain_path)
assert opt.arch == pretrain['arch']
model.load_state_dict(pretrain['state_dict'])
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
return model, parameters
return model, model.parameters()
def generate_model(opt):
assert opt.model in [
'resnet', 'preresnet', 'wideresnet', 'resnext', 'densenet'
]
if opt.model == 'resnet':
assert opt.model_depth in [10, 18, 34, 50, 101, 152, 200]
from models.resnet import get_fine_tuning_parameters
if opt.model_depth == 10:
model = resnet.resnet10(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 18:
model = resnet.resnet18(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 34:
model = resnet.resnet34(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 50:
model = resnet.resnet50(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 101:
model = resnet.resnet101(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 152:
model = resnet.resnet152(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 200:
model = resnet.resnet200(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'wideresnet':
assert opt.model_depth in [50]
from models.wide_resnet import get_fine_tuning_parameters
if opt.model_depth == 50:
model = wide_resnet.resnet50(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
k=opt.wide_resnet_k,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'resnext':
assert opt.model_depth in [50, 101, 152]
from models.resnext import get_fine_tuning_parameters
if opt.model_depth == 50:
model = resnext.resnet50(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 101:
model = resnext.resnet101(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 152:
model = resnext.resnet152(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
cardinality=opt.resnext_cardinality,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'preresnet':
assert opt.model_depth in [18, 34, 50, 101, 152, 200]
from models.pre_act_resnet import get_fine_tuning_parameters
if opt.model_depth == 18:
model = pre_act_resnet.resnet18(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 34:
model = pre_act_resnet.resnet34(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 50:
model = pre_act_resnet.resnet50(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 101:
model = pre_act_resnet.resnet101(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 152:
model = pre_act_resnet.resnet152(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 200:
model = pre_act_resnet.resnet200(
num_classes=opt.n_classes,
shortcut_type=opt.resnet_shortcut,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model == 'densenet':
assert opt.model_depth in [121, 169, 201, 264]
from models.densenet import get_fine_tuning_parameters
if opt.model_depth == 121:
model = densenet.densenet121(
num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 169:
model = densenet.densenet169(
num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 201:
model = densenet.densenet201(
num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
elif opt.model_depth == 264:
model = densenet.densenet264(
num_classes=opt.n_classes,
sample_size=opt.sample_size,
sample_duration=opt.sample_duration)
if not opt.no_cuda:
model = model.cuda()
model = nn.DataParallel(model, device_ids=None)
if opt.pretrain_path:
print('loading pretrained model {}'.format(opt.pretrain_path))
pretrain = torch.load(opt.pretrain_path)
assert opt.arch == pretrain['arch']
model.load_state_dict(pretrain['state_dict'])
if opt.model == 'densenet':
model.module.classifier = nn.Linear(
model.module.classifier.in_features, opt.n_finetune_classes)
model.module.classifier = model.module.classifier.cuda()
else:
model.module.fc = nn.Linear(model.module.fc.in_features,
opt.n_finetune_classes)
model.module.fc = model.module.fc.cuda()
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
return model, parameters
else:
if opt.pretrain_path:
print('loading pretrained model {}'.format(opt.pretrain_path))
pretrain = torch.load(opt.pretrain_path)
assert opt.arch == pretrain['arch']
model.load_state_dict(pretrain['state_dict'])
if opt.model == 'densenet':
model.classifier = nn.Linear(
model.classifier.in_features, opt.n_finetune_classes)
else:
model.fc = nn.Linear(model.fc.in_features,
opt.n_finetune_classes)
parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
return model, parameters
return model, model.parameters()
