julius/prototorch_models
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Use 'num_' in all variable names

Browse Source
This commit is contained in:
Alexander Engelsberger
2021-05-25 15:41:10 +02:00
parent e7e6bf9173
commit 72e064338c
8 changed files with 24 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
prototorch/models/abstract.py
View File

@@ -31,9 +31,9 @@ class PrototypeImageModel(pl.LightningModule):
def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
self.proto_layer.components.data.clamp_(0.0, 1.0)
def get_prototype_grid(self, nrow=2, return_channels_last=True):
def get_prototype_grid(self, num_columns=2, return_channels_last=True):
from torchvision.utils import make_grid
grid = make_grid(self.components, nrow=nrow)
grid = make_grid(self.components, nrow=num_columns)
if return_channels_last:
grid = grid.permute((1, 2, 0))
return grid.cpu()