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

Add get_prototype_grid method

Browse Source
This commit is contained in:
Jensun Ravichandran 2021-05-21 13:11:48 +02:00
parent 0933a88a1b
commit a9382dcd9b
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
prototorch/models/abstract.py
View File

@ -31,6 +31,13 @@ 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):
from torchvision.utils import make_grid
grid = make_grid(self.components, nrow=nrow)
if return_channels_last:
grid = grid.permute((1, 2, 0))
return grid.cpu()
class SiamesePrototypeModel(pl.LightningModule):
def configure_optimizers(self):