def load_defined_model(name, num_classes):
model = models.__dict__[name](num_classes=num_classes)
#Densenets don't (yet) pass on num_classes, hack it in for 169
if name == 'densenet169':
model = torchvision.models.DenseNet(num_init_features=64, growth_rate=32, \
block_config=(6, 12, 32, 32), num_classes=num_classes)
pretrained_state = model_zoo.load_url(model_urls[name])
#Diff
diff = [s for s in diff_states(model.state_dict(), pretrained_state)]
print("Replacing the following state from initialized", name, ":", \
[d[0] for d in diff])
for name, value in diff:
pretrained_state[name] = value
assert len([s for s in diff_states(model.state_dict(), pretrained_state)]) == 0
#Merge
model.load_state_dict(pretrained_state)
return model, diff
python类models()的实例源码
def build_cnn(args, dtype):
if not hasattr(torchvision.models, args.cnn_model):
raise ValueError('Invalid model "%s"' % args.cnn_model)
if not 'resnet' in args.cnn_model:
raise ValueError('Feature extraction only supports ResNets')
whole_cnn = getattr(torchvision.models, args.cnn_model)(pretrained=True)
layers = [
whole_cnn.conv1,
whole_cnn.bn1,
whole_cnn.relu,
whole_cnn.maxpool,
]
for i in range(args.cnn_model_stage):
name = 'layer%d' % (i + 1)
layers.append(getattr(whole_cnn, name))
cnn = torch.nn.Sequential(*layers)
cnn.type(dtype)
cnn.eval()
return cnn
def build_model(args):
if not hasattr(torchvision.models, args.model):
raise ValueError('Invalid model "%s"' % args.model)
if not 'resnet' in args.model:
raise ValueError('Feature extraction only supports ResNets')
cnn = getattr(torchvision.models, args.model)(pretrained=True)
layers = [
cnn.conv1,
cnn.bn1,
cnn.relu,
cnn.maxpool,
]
for i in range(args.model_stage):
name = 'layer%d' % (i + 1)
layers.append(getattr(cnn, name))
model = torch.nn.Sequential(*layers)
model.cuda()
model.eval()
return model
def load_imagenet_model(name, *args, **kwargs):
# load model and state dict
model = getattr(torchvision.models, name)(*args, **kwargs)
state_dict = torch.utils.model_zoo.load_url(model_path[name])
# load state dict to model
load_state_dict(model, state_dict)
return model
def _construct_flow_model(self, base_model):
# modify the convolution layers
# Torch models are usually defined in a hierarchical way.
# nn.modules.children() return all sub modules in a DFS manner
modules = list(self.base_model.modules())
first_conv_idx = list(filter(lambda x: isinstance(modules[x], nn.Conv2d), list(range(len(modules)))))[0]
conv_layer = modules[first_conv_idx]
container = modules[first_conv_idx - 1]
# modify parameters, assume the first blob contains the convolution kernels
params = [x.clone() for x in conv_layer.parameters()]
kernel_size = params[0].size()
new_kernel_size = kernel_size[:1] + (2 * self.new_length,) + kernel_size[2:]
new_kernels = params[0].data.mean(dim=1, keepdim=True).expand(new_kernel_size).contiguous()
new_conv = nn.Conv2d(2 * self.new_length, conv_layer.out_channels,
conv_layer.kernel_size, conv_layer.stride, conv_layer.padding,
bias=True if len(params) == 2 else False)
new_conv.weight.data = new_kernels
if len(params) == 2:
new_conv.bias.data = params[1].data # add bias if neccessary
layer_name = list(container.state_dict().keys())[0][:-7] # remove .weight suffix to get the layer name
# replace the first convlution layer
setattr(container, layer_name, new_conv)
return base_model
def _construct_diff_model(self, base_model, keep_rgb=False):
# modify the convolution layers
# Torch models are usually defined in a hierarchical way.
# nn.modules.children() return all sub modules in a DFS manner
modules = list(self.base_model.modules())
first_conv_idx = filter(lambda x: isinstance(modules[x], nn.Conv2d), list(range(len(modules))))[0]
conv_layer = modules[first_conv_idx]
container = modules[first_conv_idx - 1]
# modify parameters, assume the first blob contains the convolution kernels
params = [x.clone() for x in conv_layer.parameters()]
kernel_size = params[0].size()
if not keep_rgb:
new_kernel_size = kernel_size[:1] + (3 * self.new_length,) + kernel_size[2:]
new_kernels = params[0].data.mean(dim=1, keepdim=True).expand(new_kernel_size).contiguous()
else:
new_kernel_size = kernel_size[:1] + (3 * self.new_length,) + kernel_size[2:]
new_kernels = torch.cat((params[0].data, params[0].data.mean(dim=1, keepdim=True).expand(new_kernel_size).contiguous()),
1)
new_kernel_size = kernel_size[:1] + (3 + 3 * self.new_length,) + kernel_size[2:]
new_conv = nn.Conv2d(new_kernel_size[1], conv_layer.out_channels,
conv_layer.kernel_size, conv_layer.stride, conv_layer.padding,
bias=True if len(params) == 2 else False)
new_conv.weight.data = new_kernels
if len(params) == 2:
new_conv.bias.data = params[1].data # add bias if neccessary
layer_name = list(container.state_dict().keys())[0][:-7] # remove .weight suffix to get the layer name
# replace the first convolution layer
setattr(container, layer_name, new_conv)
return base_model
def load_weights(self, base_file):
other, ext = os.path.splitext(base_file)
if ext == '.pkl' or '.pth':
print('Loading weights into state dict...')
self.load_state_dict(torch.load(base_file))
print('Finished!')
else:
print('Error: Sorry Only .pth and .pkl files currently supported!')
# This function is derived from torchvision VGG make_layers()
# https://github.com/pytorch/vision/blob/master/torchvision/models/vgg.py
def load_weights(self, base_file):
other, ext = os.path.splitext(base_file)
if ext == '.pkl' or '.pth':
print('Loading weights into state dict...')
self.load_state_dict(torch.load(base_file))
print('Finished!')
else:
print('Sorry only .pth and .pkl files supported.')
# This function is derived from torchvision VGG make_layers()
# https://github.com/pytorch/vision/blob/master/torchvision/models/vgg.py
def __init__(self, model_name):
if model_name in MODEL_NAMES:
model_name = MODEL_NAMES[model_name]
else:
assert model_name in MODEL_NAMES.values(), MODEL_NAMES
self.net = getattr(models, model_name)(pretrained=True)
self.labels = self.process_labels()
self.normalize = Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
def init_cnn(self, opt):
"""Lazy initialization of preprocessor model in case we don't need any
image preprocessing.
"""
try:
import torch
self.use_cuda = (not opt.get('no_cuda', False)
and torch.cuda.is_available())
self.torch = torch
except ModuleNotFoundError:
raise ModuleNotFoundError('Need to install Pytorch: go to pytorch.org')
from torch.autograd import Variable
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
try:
import h5py
self.h5py = h5py
except ModuleNotFoundError:
raise ModuleNotFoundError('Need to install h5py')
if 'image_mode' not in opt or 'image_size' not in opt:
raise RuntimeError(
'Need to add image arguments to opt. See '
'parlai.core.params.ParlaiParser.add_image_args')
self.image_mode = opt['image_mode']
self.image_size = opt['image_size']
self.crop_size = opt['image_cropsize']
if self.use_cuda:
print('[ Using CUDA ]')
torch.cuda.set_device(opt.get('gpu', -1))
cnn_type, layer_num = self.image_mode_switcher()
# initialize the pretrained CNN using pytorch.
CNN = getattr(torchvision.models, cnn_type)
# cut off the additional layer.
self.netCNN = nn.Sequential(
*list(CNN(pretrained=True).children())[:layer_num])
# initialize the transform function using torch vision.
self.transform = transforms.Compose([
transforms.Scale(self.image_size),
transforms.CenterCrop(self.crop_size),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
# container for single image
self.xs = torch.zeros(1, 3, self.crop_size, self.crop_size)
if self.use_cuda:
self.netCNN.cuda()
self.xs = self.xs.cuda()
# make self.xs variable.
self.xs = Variable(self.xs)
def _prepare_base_model(self, base_model):
if 'resnet' in base_model or 'vgg' in base_model:
self.base_model = getattr(torchvision.models, base_model)(True)
self.base_model.last_layer_name = 'fc'
self.input_size = 224
self.input_mean = [0.485, 0.456, 0.406]
self.input_std = [0.229, 0.224, 0.225]
if self.modality == 'Flow':
self.input_mean = [0.5]
self.input_std = [np.mean(self.input_std)]
elif self.modality == 'RGBDiff':
self.input_mean = [0.485, 0.456, 0.406] + [0] * 3 * self.new_length
self.input_std = self.input_std + [np.mean(self.input_std) * 2] * 3 * self.new_length
elif base_model == 'BNInception':
import model_zoo
self.base_model = getattr(model_zoo, base_model)()
self.base_model.last_layer_name = 'fc'
self.input_size = 224
self.input_mean = [104, 117, 128]
self.input_std = [1]
if self.modality == 'Flow':
self.input_mean = [128]
elif self.modality == 'RGBDiff':
self.input_mean = self.input_mean * (1 + self.new_length)
elif base_model == 'InceptionV3':
import model_zoo
self.base_model = getattr(model_zoo, base_model)()
self.base_model.last_layer_name = 'top_cls_fc'
self.input_size = 299
self.input_mean = [104, 117, 128]
self.input_std = [1]
if self.modality == 'Flow':
self.input_mean = [128]
elif self.modality == 'RGBDiff':
self.input_mean = self.input_mean * (1 + self.new_length)
elif 'inception' in base_model:
import model_zoo
self.base_model = getattr(model_zoo, base_model)()
self.base_model.last_layer_name = 'classif'
self.input_size = 299
self.input_mean = [0.5]
self.input_std = [0.5]
else:
raise ValueError('Unknown base model: {}'.format(base_model))
