luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

curvature training

Browse Source
This commit is contained in:
luni
2023-04-08 00:43:08 +02:00
parent 52e5658dcc
commit cada175296
7 changed files with 309 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/base.py
+5 -5
View File
@@ -33,10 +33,10 @@ class VariationalEncoder(nn.Module):
class AEEncoder(nn.Module):
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', device='cuda'):
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda'):
super(AEEncoder, self).__init__()
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device)
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device)
if activation == 'tanh':
self.act = nn.Tanh()
elif activation == 'sigmoid':
@@ -81,10 +81,10 @@ class Decoder(nn.Module):
x = self.decoder(x)
if self.is_label:
x = x * mask[:,:,:,None]
x = x * mask[:,None,:,:]
background = torch.ones((mask.shape[0],512,256,1), device=mask.get_device()) * (1 - mask[:,:,:,None])
x = torch.concat((background, x), axis=-1)
background = torch.ones((mask.shape[0],1,512,256), device=mask.get_device()) * (1 - mask[:,None,:,:])
x = torch.concat((background, x), dim=1)
x = self.softmax(x)
if len(inputs) == 2:
models/cond_encoder.py
+2 -2
View File
@@ -8,8 +8,8 @@ class AutoEncoder(nn.Module):
def __init__(self, config):
super(AutoEncoder, self).__init__()
self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], device=config['device'])
self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device'])
self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], is_label=config['is_label'], device=config['device'], nr_of_labels=7)
self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device'], nr_of_classes=6)
def forward(self, input):
models/convnext_orig.py
+4 -2
View File
@@ -64,6 +64,8 @@ class ConvNeXtBase(nn.Module):
norm_layer: Optional[Callable[..., nn.Module]] = None,
device: str = 'cuda',
kernel_size: int = 7,
is_label: bool = False,
nr_of_labels: Optional[int] = None,
) -> None:
super().__init__()
@@ -83,7 +85,7 @@ class ConvNeXtBase(nn.Module):
# Stem
firstconv_output_channels = block_setting[0].input_channels
self.stem = Conv2dNormActivation(
3,
3 if not is_label else nr_of_labels,
firstconv_output_channels,
kernel_size=4,
stride=4,
@@ -234,7 +236,7 @@ def _convnext(
if reversed:
model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, is_label=is_label, nr_of_labels=nr_of_labels, kernel_size=kernel_size, device=device, **kwargs)
else:
model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, device=device, **kwargs)
model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device, **kwargs)
return model