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: