
--- train.py
+++ train.py
... | ... | @@ -91,8 +91,8 @@ |
91 | 91 |
ARNN_loss_window = vis.line(Y=np.array([0]), X=np.array([0]), opts=dict(title='Generator-AttentionRNN Loss')) |
92 | 92 |
AE_loss_window = vis.line(Y=np.array([0]), X=np.array([0]), opts=dict(title='Generator-AutoEncoder Loss')) |
93 | 93 |
Discriminator_loss_window = vis.line(Y=np.array([0]), X=np.array([0]), opts=dict(title='Discriminator Loss')) |
94 |
-attenton_map_visualizer = vis.image(np.zeros((692,776)), opts=dict(title='Attention Map')) |
|
95 |
- |
|
94 |
+Attention_map_visualizer = vis.image(np.zeros((692, 776)), opts=dict(title='Attention Map')) |
|
95 |
+Generator_output_visualizer = vis.image(np.zeros((692,776)), opts=dict(title='Generated Derain Output')) |
|
96 | 96 |
|
97 | 97 |
for epoch_num, epoch in enumerate(range(epochs)): |
98 | 98 |
for i, imgs in enumerate(dataloader): |
... | ... | @@ -146,7 +146,6 @@ |
146 | 146 |
} |
147 | 147 |
|
148 | 148 |
logger.print_training_log(epoch_num, epochs, i, len(dataloader), losses) |
149 |
- |
|
150 | 149 |
# visdom logger |
151 | 150 |
vis.line(Y=np.array([generator_loss_ARNN.item()]), X=np.array([epoch * epoch_num + i]), win=ARNN_loss_window, |
152 | 151 |
update='append') |
... | ... | @@ -154,8 +153,8 @@ |
154 | 153 |
update='append') |
155 | 154 |
vis.line(Y=np.array([discriminator_loss.item()]), X=np.array([epoch * epoch_num + i]), win=Discriminator_loss_window, |
156 | 155 |
update='append') |
157 |
- vis.image(generator_attention_map[-1], win=attenton_map_visualizer, opts=dict(title="Attention Map")) |
|
158 |
- |
|
156 |
+ vis.image(generator_attention_map[-1][0,0,:,:], win=Attention_map_visualizer, opts=dict(title="Attention Map")) |
|
157 |
+ vis.image(generator_result['skip_3'][-1][0,0,:,:], win=Generator_output_visualizer, opts=dict(title="Generator Output")) |
|
159 | 158 |
day = strftime("%Y-%m-%d %H:%M:%S", gmtime()) |
160 | 159 |
if epoch % save_interval == 0 and epoch != 0: |
161 | 160 |
torch.save(generator.attentiveRNN.state_dict(), f"weight/Attention_RNN_{epoch}_{day}.pt") |
Add a comment
Delete comment
Once you delete this comment, you won't be able to recover it. Are you sure you want to delete this comment?