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

Fix attention images

Browse Source
This commit is contained in:
James Betker 2020-10-09 19:21:55 -06:00
parent 7e777ea34c
commit fe50d6f9d0
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -557,10 +557,10 @@ class SwitchModelBase(nn.Module):
temp = max(1, 1 + self.init_temperature * temp = max(1, 1 + self.init_temperature *
(self.final_temperature_step - step) / self.final_temperature_step) (self.final_temperature_step - step) / self.final_temperature_step)
self.set_temperature(temp) self.set_temperature(temp)
if step % 200 == 0: if step % 100 == 0:
output_path = os.path.join(experiments_path, "attention_maps") output_path = os.path.join(experiments_path, "attention_maps")
prefix = "amap_%i_a%i_%%i.png" prefix = "amap_%i_a%i_%%i.png"
[save_attention_to_image_rgb(output_path, self.attentions[i], self.nf, prefix % (step, i), step, [save_attention_to_image_rgb(output_path, self.attentions[i], self.attentions[i].shape[3], prefix % (step, i), step,
output_mag=False) for i in range(len(self.attentions))] output_mag=False) for i in range(len(self.attentions))]
if self.lr is not None: if self.lr is not None:
torchvision.utils.save_image(self.lr[:, :3], os.path.join(experiments_path, "attention_maps", torchvision.utils.save_image(self.lr[:, :3], os.path.join(experiments_path, "attention_maps",