Skip to content

GitLab

"git@developer.sourcefind.cn:OpenDAS/ollama.git" did not exist on "ac0801ecedec76492a8a024ad2f4fb84f80f001c"
Commit 47bc8153 authored by Konstantin Lopuhin's avatar Konstantin Lopuhin Committed by Soumith Chintala
Browse files

pass kwargs to densenet (#199)

parent 08b1f59f
Show whitespace changes
Inline Side-by-side
Showing with 8 additions and 4 deletions
torchvision/models/densenet.py
View file @ 47bc8153
...@@ -22,7 +22,8 @@ def densenet121(pretrained=False, **kwargs): ...@@ -22,7 +22,8 @@ def densenet121(pretrained=False, **kwargs):
Args: Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet pretrained (bool): If True, returns a model pre-trained on ImageNet
""" """
model = DenseNet(num_init_features=64, growth_rate=32, block_config=(6, 12, 24, 16)) model = DenseNet(num_init_features=64, growth_rate=32, block_config=(6, 12, 24, 16),
**kwargs)
if pretrained: if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['densenet121'])) model.load_state_dict(model_zoo.load_url(model_urls['densenet121']))
return model return model
...@@ -35,7 +36,8 @@ def densenet169(pretrained=False, **kwargs): ...@@ -35,7 +36,8 @@ def densenet169(pretrained=False, **kwargs):
Args: Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet pretrained (bool): If True, returns a model pre-trained on ImageNet
""" """
model = DenseNet(num_init_features=64, growth_rate=32, block_config=(6, 12, 32, 32)) model = DenseNet(num_init_features=64, growth_rate=32, block_config=(6, 12, 32, 32),
**kwargs)
if pretrained: if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['densenet169'])) model.load_state_dict(model_zoo.load_url(model_urls['densenet169']))
return model return model
...@@ -48,7 +50,8 @@ def densenet201(pretrained=False, **kwargs): ...@@ -48,7 +50,8 @@ def densenet201(pretrained=False, **kwargs):
Args: Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet pretrained (bool): If True, returns a model pre-trained on ImageNet
""" """
model = DenseNet(num_init_features=64, growth_rate=32, block_config=(6, 12, 48, 32)) model = DenseNet(num_init_features=64, growth_rate=32, block_config=(6, 12, 48, 32),
**kwargs)
if pretrained: if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['densenet201'])) model.load_state_dict(model_zoo.load_url(model_urls['densenet201']))
return model return model
...@@ -61,7 +64,8 @@ def densenet161(pretrained=False, **kwargs): ...@@ -61,7 +64,8 @@ def densenet161(pretrained=False, **kwargs):
Args: Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet pretrained (bool): If True, returns a model pre-trained on ImageNet
""" """
model = DenseNet(num_init_features=96, growth_rate=48, block_config=(6, 12, 36, 24)) model = DenseNet(num_init_features=96, growth_rate=48, block_config=(6, 12, 36, 24),
**kwargs)
if pretrained: if pretrained:
model.load_state_dict(model_zoo.load_url(model_urls['densenet161'])) model.load_state_dict(model_zoo.load_url(model_urls['densenet161']))
return model return model
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment