mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Spinenet should allow bypassing the initial conv

Browse Source 
This makes feeding in references for recurrence easier.
This commit is contained in:
James Betker 2020-10-17 20:16:47 -06:00
parent c7f3fc4dd9
commit b008a27d39
3 changed files with 6 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
codes/data/chunk_with_reference.py
View File

@ -32,7 +32,7 @@ class ChunkWithReference:
elif self.strict:
raise FileNotFoundError(tile_id, self.tiles[item])
else:
center = torch.tensor([128,128], dtype=torch.long)
center = torch.tensor([128, 128], dtype=torch.long)
tile_width = 256
mask = np.full(tile.shape[:2] + (1,), fill_value=.1, dtype=tile.dtype)
mask[center[0] - tile_width // 2:center[0] + tile_width // 2, center[1] - tile_width // 2:center[1] + tile_width // 2] = 1

4
codes/models/archs/ChainedEmbeddingGen.py
View File

@ -58,8 +58,8 @@ class ChainedEmbeddingGen(nn.Module):
self.upsample = FinalUpsampleBlock2x(64)
def forward(self, x):
emb = checkpoint(self.spine, x)
fea = self.initial_conv(x)
emb = checkpoint(self.spine, fea)
for block in self.blocks:
fea = fea + checkpoint(block, fea, *emb)
return checkpoint(self.upsample, fea),
@ -82,11 +82,11 @@ class ChainedEmbeddingGenWithStructure(nn.Module):
self.upsample = FinalUpsampleBlock2x(64)
def forward(self, x, recurrent=None):
emb = checkpoint(self.spine, x)
fea = self.initial_conv(x)
if self.recurrent:
rec = self.recurrent_process(recurrent)
fea, _ = self.recurrent_join(fea, rec)
emb = checkpoint(self.spine, fea)
grad = fea
for i, block in enumerate(self.blocks):
fea = fea + checkpoint(block, fea, *emb)

11
codes/models/archs/spinenet_arch.py
View File

@ -245,12 +245,7 @@ class SpineNet(nn.Module):
stride=2)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
else:
self.conv1 = ConvGnSilu(
in_channels,
64,
kernel_size=7,
stride=1)
self.maxpool = None
self.conv1 = None
# Build the initial level 2 blocks.
self.init_block1 = make_res_layer(
@ -311,8 +306,8 @@ class SpineNet(nn.Module):
std = torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(input.device)
input = (input - mean) / std
feat = self.conv1(input)
if self.maxpool:
if self.conv1 is not None:
feat = self.conv1(input)
feat = self.maxpool(feat)
feat1 = self.init_block1(feat)
feat2 = self.init_block2(feat1)