Update training script to use weight normalization

Coloquinte · Jun 22, 2021 · db7f15e · db7f15e
1 parent 2ce78e0
commit db7f15e
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/scripts/train_ninasr.sh b/scripts/train_ninasr.sh
@@ -11,7 +11,8 @@ train_network () {
     python -m torchsr.train --arch $arch --scale $scale --tune-backend \
         --log-dir logs_train/${arch}_x${scale} --save-checkpoint ${arch}_x${scale}.pt \
         --lr $learning_rate --epochs $epochs --lr-decay-steps $((epochs*2/3)) $((epochs*5/6)) --lr-decay-rate 3 \
-        --patch-size-train $(( (patch_size+1) * scale)) --shave-border $scale --replication-pad 4
+        --patch-size-train $(( (patch_size+1) * scale)) --shave-border $scale --replication-pad 4 \
+        --weight-norm
 }
 
 train_network_with_pretrained () {
@@ -39,7 +40,8 @@ train_network_with_pretrained () {
         --log-dir logs_train/${arch}_x${scale} --save-checkpoint ${arch}_x${scale}.pt \
         --load-pretrained ${arch}_x${scale}_pre_model.pt \
         --lr $learning_rate --epochs $epochs --lr-decay-steps $((epochs*2/3)) $((epochs*5/6)) --lr-decay-rate 3 \
-        --patch-size-train $(( (patch_size+1) * scale)) --shave-border $scale --replication-pad 4
+        --patch-size-train $(( (patch_size+1) * scale)) --shave-border $scale --replication-pad 4 \
+        --weight-norm
 }
 
 epochs=300
@@ -48,10 +50,9 @@ learning_rate=0.001
 
 # NinaSR-B0
 train_network ninasr_b0 2
-train_network_with_pretrained ninasr_b0 3 2
-train_network_with_pretrained ninasr_b0 4 3
-train_network_with_pretrained ninasr_b0 8 4
-train_network_with_pretrained ninasr_b0 2 4
+train_network ninasr_b0 3 
+train_network ninasr_b0 4 
+train_network ninasr_b0 8 
 
 epochs=500
 patch_size=48