我们从Python开源项目中，提取了以下11个代码示例，用于说明如何使用torchvision.models()。
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
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))
