Update train.py

train.py

@@ -181,63 +181,65 @@ 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
 
-    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)
 
 if __name__ == "__main__":
+    print(tf.config.list_physical_devices('GPU'))
     opt = parse_opt()
     main(opt.content_img_path, opt.style_img_path, opt.epochs, opt.step_per_epoch, opt.learning_rate, opt.content_loss_factor, opt.style_loss_factor, opt.img_size, opt.img_width, opt.img_height)