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

direct training

Browse Source
This commit is contained in:
luni
2023-05-23 22:57:08 +02:00
parent 843c9ac4b0
commit 53217c6e5d
8 changed files with 470 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/base.py
+2 -2
View File
@@ -35,10 +35,10 @@ class VariationalEncoder(nn.Module):
class AEEncoder(nn.Module):
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda'):
def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda', input_features=3):
super(AEEncoder, self).__init__()
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, 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, input_features=input_features)
if activation == 'tanh':
self.act = nn.Tanh()
elif activation == 'sigmoid':