fromtypingimportUnion,List,Dict,Any,castimporttorchimporttorch.nnasnnfrom.._internally_replaced_utilsimportload_state_dict_from_urlfrom..utilsimport_log_api_usage_once__all__=["VGG","vgg11","vgg11_bn","vgg13","vgg13_bn","vgg16","vgg16_bn","vgg19_bn","vgg19",]model_urls={"vgg11":"https://download.pytorch.org/models/vgg11-8a719046.pth","vgg13":"https://download.pytorch.org/models/vgg13-19584684.pth","vgg16":"https://download.pytorch.org/models/vgg16-397923af.pth","vgg19":"https://download.pytorch.org/models/vgg19-dcbb9e9d.pth","vgg11_bn":"https://download.pytorch.org/models/vgg11_bn-6002323d.pth","vgg13_bn":"https://download.pytorch.org/models/vgg13_bn-abd245e5.pth","vgg16_bn":"https://download.pytorch.org/models/vgg16_bn-6c64b313.pth","vgg19_bn":"https://download.pytorch.org/models/vgg19_bn-c79401a0.pth",}classVGG(nn.Module):def__init__(self,features:nn.Module,num_classes:int=1000,init_weights:bool=True,dropout:float=0.5)->None:super().__init__()_log_api_usage_once(self)self.features=featuresself.avgpool=nn.AdaptiveAvgPool2d((7,7))self.classifier=nn.Sequential(nn.Linear(512*7*7,4096),nn.ReLU(True),nn.Dropout(p=dropout),nn.Linear(4096,4096),nn.ReLU(True),nn.Dropout(p=dropout),nn.Linear(4096,num_classes),)ifinit_weights:forminself.modules():ifisinstance(m,nn.Conv2d):nn.init.kaiming_normal_(m.weight,mode="fan_out",nonlinearity="relu")ifm.biasisnotNone:nn.init.constant_(m.bias,0)elifisinstance(m,nn.BatchNorm2d):nn.init.constant_(m.weight,1)nn.init.constant_(m.bias,0)elifisinstance(m,nn.Linear):nn.init.normal_(m.weight,0,0.01)nn.init.constant_(m.bias,0)defforward(self,x:torch.Tensor)->torch.Tensor:x=self.features(x)x=self.avgpool(x)x=torch.flatten(x,1)x=self.classifier(x)returnxdefmake_layers(cfg:List[Union[str,int]],batch_norm:bool=False)->nn.Sequential:layers:List[nn.Module]=[]in_channels=3forvincfg:ifv=="M":layers+=[nn.MaxPool2d(kernel_size=2,stride=2)]else:v=cast(int,v)conv2d=nn.Conv2d(in_channels,v,kernel_size=3,padding=1)ifbatch_norm:layers+=[conv2d,nn.BatchNorm2d(v),nn.ReLU(inplace=True)]else:layers+=[conv2d,nn.ReLU(inplace=True)]in_channels=vreturnnn.Sequential(*layers)cfgs:Dict[str,List[Union[str,int]]]={"A":[64,"M",128,"M",256,256,"M",512,512,"M",512,512,"M"],"B":[64,64,"M",128,128,"M",256,256,"M",512,512,"M",512,512,"M"],"D":[64,64,"M",128,128,"M",256,256,256,"M",512,512,512,"M",512,512,512,"M"],"E":[64,64,"M",128,128,"M",256,256,256,256,"M",512,512,512,512,"M",512,512,512,512,"M"],}def_vgg(arch:str,cfg:str,batch_norm:bool,pretrained:bool,progress:bool,**kwargs:Any)->VGG:ifpretrained:kwargs["init_weights"]=Falsemodel=VGG(make_layers(cfgs[cfg],batch_norm=batch_norm),**kwargs)ifpretrained:state_dict=load_state_dict_from_url(model_urls[arch],progress=progress)model.load_state_dict(state_dict)returnmodeldefvgg11(pretrained:bool=False,progress:bool=True,**kwargs:Any)->VGG:r"""VGG 11-layer model (configuration "A") from `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_. The required minimum input size of the model is 32x32. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """return_vgg("vgg11","A",False,pretrained,progress,**kwargs)defvgg11_bn(pretrained:bool=False,progress:bool=True,**kwargs:Any)->VGG:r"""VGG 11-layer model (configuration "A") with batch normalization `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_. The required minimum input size of the model is 32x32. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """return_vgg("vgg11_bn","A",True,pretrained,progress,**kwargs)defvgg13(pretrained:bool=False,progress:bool=True,**kwargs:Any)->VGG:r"""VGG 13-layer model (configuration "B") `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_. The required minimum input size of the model is 32x32. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """return_vgg("vgg13","B",False,pretrained,progress,**kwargs)defvgg13_bn(pretrained:bool=False,progress:bool=True,**kwargs:Any)->VGG:r"""VGG 13-layer model (configuration "B") with batch normalization `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_. The required minimum input size of the model is 32x32. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """return_vgg("vgg13_bn","B",True,pretrained,progress,**kwargs)
[docs]defvgg16(pretrained:bool=False,progress:bool=True,**kwargs:Any)->VGG:r"""VGG 16-layer model (configuration "D") `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_. The required minimum input size of the model is 32x32. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """return_vgg("vgg16","D",False,pretrained,progress,**kwargs)
[docs]defvgg16_bn(pretrained:bool=False,progress:bool=True,**kwargs:Any)->VGG:r"""VGG 16-layer model (configuration "D") with batch normalization `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_. The required minimum input size of the model is 32x32. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """return_vgg("vgg16_bn","D",True,pretrained,progress,**kwargs)
[docs]defvgg19(pretrained:bool=False,progress:bool=True,**kwargs:Any)->VGG:r"""VGG 19-layer model (configuration "E") `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_. The required minimum input size of the model is 32x32. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """return_vgg("vgg19","E",False,pretrained,progress,**kwargs)
[docs]defvgg19_bn(pretrained:bool=False,progress:bool=True,**kwargs:Any)->VGG:r"""VGG 19-layer model (configuration 'E') with batch normalization `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_. The required minimum input size of the model is 32x32. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet progress (bool): If True, displays a progress bar of the download to stderr """return_vgg("vgg19_bn","E",True,pretrained,progress,**kwargs)
Docs
Access comprehensive developer documentation for PyTorch
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.