Update train.py

train.py

@@ -179,61 +179,61 @@ def train_one_step(model, noise_image, optimizer, target_content_features, targe
 def main(content_img, style_img, epochs, step_per_epoch, learning_rate, content_loss_factor, style_loss_factor, img_size, img_width, img_height):
     global CONTENT_LOSS_FACTOR, STYLE_LOSS_FACTOR, CONTENT_IMAGE_PATH, STYLE_IMAGE_PATH, OUTPUT_DIR, EPOCHS, LEARNING_RATE, STEPS_PER_EPOCH, M, N, image_mean, image_std, IMG_WIDTH, IMG_HEIGHT
 
-    with tf.device('/cuda:0'):
-        CONTENT_LOSS_FACTOR = content_loss_factor
-        STYLE_LOSS_FACTOR = style_loss_factor
-        CONTENT_IMAGE_PATH = content_img
-        STYLE_IMAGE_PATH = style_img
-        EPOCHS = epochs
-        LEARNING_RATE = learning_rate
-        STEPS_PER_EPOCH = step_per_epoch
-    
-        # 内容特征层及损失加权系数
-        CONTENT_LAYERS = {"block4_conv2": 0.5, "block5_conv2": 0.5}
-        # 风格特征层及损失加权系数
-        STYLE_LAYERS = {
-            "block1_conv1": 0.2,
-            "block2_conv1": 0.2,
-            "block3_conv1": 0.2,
-            "block4_conv1": 0.2,
-            "block5_conv1": 0.2,
-        }
-    
-        if img_size == "default size":
-            IMG_WIDTH = 450
-            IMG_HEIGHT = 300
-        else:
-            IMG_WIDTH = img_width
-            IMG_HEIGHT = img_height
-    
-        print("IMG_WIDTH:", IMG_WIDTH)
-        print("IMG_HEIGHT:", IMG_HEIGHT)
-    
-        # 我们准备使用经典网络在imagenet数据集上的预训练权重，所以归一化时也要使用imagenet的平均值和标准差
-        image_mean = tf.constant([0.485, 0.456, 0.406])
-        image_std = tf.constant([0.299, 0.224, 0.225])
-    
-        model = NeuralStyleTransferModel(CONTENT_LAYERS, STYLE_LAYERS)
-    
-        content_image = load_images_from_list(CONTENT_IMAGE_PATH, IMG_WIDTH, IMG_HEIGHT)
-        style_image = load_images_from_list(STYLE_IMAGE_PATH, IMG_WIDTH, IMG_HEIGHT)
-    
-        target_content_features = model(content_image)["content"]
-        target_style_features = model(style_image)["style"]
-    
-        M = IMG_WIDTH * IMG_HEIGHT
-        N = 3
-    
-        optimizer = tf.keras.optimizers.Adam(LEARNING_RATE)
-    
-        noise_image = tf.Variable((content_image[0] + np.random.uniform(-0.2, 0.2, (1, IMG_HEIGHT, IMG_WIDTH, 3))) / 2)
-    
-        for epoch in range(EPOCHS):
-            with tqdm(total=STEPS_PER_EPOCH, desc="Epoch {}/{}".format(epoch + 1, EPOCHS)) as pbar:
-                for step in range(STEPS_PER_EPOCH):
-                    _loss = train_one_step(model, noise_image, optimizer, target_content_features, target_style_features)
-                    pbar.set_postfix({"loss": "%.4f" % float(_loss)})
-                    pbar.update(1)
+    # with tf.device('/cuda:0'):
+    CONTENT_LOSS_FACTOR = content_loss_factor
+    STYLE_LOSS_FACTOR = style_loss_factor
+    CONTENT_IMAGE_PATH = content_img
+    STYLE_IMAGE_PATH = style_img
+    EPOCHS = epochs
+    LEARNING_RATE = learning_rate
+    STEPS_PER_EPOCH = step_per_epoch
+
+    # 内容特征层及损失加权系数
+    CONTENT_LAYERS = {"block4_conv2": 0.5, "block5_conv2": 0.5}
+    # 风格特征层及损失加权系数
+    STYLE_LAYERS = {
+        "block1_conv1": 0.2,
+        "block2_conv1": 0.2,
+        "block3_conv1": 0.2,
+        "block4_conv1": 0.2,
+        "block5_conv1": 0.2,
+    }
+
+    if img_size == "default size":
+        IMG_WIDTH = 450
+        IMG_HEIGHT = 300
+    else:
+        IMG_WIDTH = img_width
+        IMG_HEIGHT = img_height
+
+    print("IMG_WIDTH:", IMG_WIDTH)
+    print("IMG_HEIGHT:", IMG_HEIGHT)
+
+    # 我们准备使用经典网络在imagenet数据集上的预训练权重，所以归一化时也要使用imagenet的平均值和标准差
+    image_mean = tf.constant([0.485, 0.456, 0.406])
+    image_std = tf.constant([0.299, 0.224, 0.225])
+
+    model = NeuralStyleTransferModel(CONTENT_LAYERS, STYLE_LAYERS)
+
+    content_image = load_images_from_list(CONTENT_IMAGE_PATH, IMG_WIDTH, IMG_HEIGHT)
+    style_image = load_images_from_list(STYLE_IMAGE_PATH, IMG_WIDTH, IMG_HEIGHT)
+
+    target_content_features = model(content_image)["content"]
+    target_style_features = model(style_image)["style"]
+
+    M = IMG_WIDTH * IMG_HEIGHT
+    N = 3
+
+    optimizer = tf.keras.optimizers.Adam(LEARNING_RATE)
+
+    noise_image = tf.Variable((content_image[0] + np.random.uniform(-0.2, 0.2, (1, IMG_HEIGHT, IMG_WIDTH, 3))) / 2)
+
+    for epoch in range(EPOCHS):
+        with tqdm(total=STEPS_PER_EPOCH, desc="Epoch {}/{}".format(epoch + 1, EPOCHS)) as pbar:
+            for step in range(STEPS_PER_EPOCH):
+                _loss = train_one_step(model, noise_image, optimizer, target_content_features, target_style_features)
+                pbar.set_postfix({"loss": "%.4f" % float(_loss)})
+                pbar.update(1)
     
     return save_image_for_gradio(noise_image)