autoforge 1.5.2__tar.gz → 1.5.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (40) hide show
  1. {autoforge-1.5.2 → autoforge-1.5.4}/PKG-INFO +9 -6
  2. {autoforge-1.5.2 → autoforge-1.5.4}/README.md +7 -5
  3. {autoforge-1.5.2 → autoforge-1.5.4}/pyproject.toml +3 -2
  4. {autoforge-1.5.2 → autoforge-1.5.4}/src/AutoForge.egg-info/PKG-INFO +9 -6
  5. {autoforge-1.5.2 → autoforge-1.5.4}/src/AutoForge.egg-info/SOURCES.txt +4 -1
  6. {autoforge-1.5.2 → autoforge-1.5.4}/src/AutoForge.egg-info/requires.txt +1 -0
  7. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Helper/CAdamW.py +27 -15
  8. autoforge-1.5.4/src/autoforge/Helper/Heightmaps/ChristofidesHeightMap.py +455 -0
  9. autoforge-1.5.4/src/autoforge/Helper/Heightmaps/DepthEstimateHeightMap.py +391 -0
  10. autoforge-1.5.4/src/autoforge/Helper/Heightmaps/OldHeightmap.py +154 -0
  11. autoforge-1.5.4/src/autoforge/Helper/OptimizerHelper.py +475 -0
  12. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Helper/PruningHelper.py +9 -3
  13. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Helper/filamentcolors_library.py +11 -5
  14. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Loss/LossFunctions.py +2 -2
  15. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Modules/Optimizer.py +1 -1
  16. autoforge-1.5.4/src/autoforge/__init__.py +0 -0
  17. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/auto_forge.py +46 -17
  18. autoforge-1.5.2/src/autoforge/Helper/OptimizerHelper.py +0 -985
  19. autoforge-1.5.2/src/autoforge/old_jax_files.py +0 -142
  20. {autoforge-1.5.2 → autoforge-1.5.4}/setup.cfg +0 -0
  21. {autoforge-1.5.2 → autoforge-1.5.4}/src/AutoForge.egg-info/dependency_links.txt +0 -0
  22. {autoforge-1.5.2 → autoforge-1.5.4}/src/AutoForge.egg-info/entry_points.txt +0 -0
  23. {autoforge-1.5.2 → autoforge-1.5.4}/src/AutoForge.egg-info/top_level.txt +0 -0
  24. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Gui/Views/__init__.py +0 -0
  25. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Gui/Views/backend_database.py +0 -0
  26. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Gui/Views/filament_editor.py +0 -0
  27. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Gui/__init__.py +0 -0
  28. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Gui/app.py +0 -0
  29. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Gui/main.py +0 -0
  30. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Helper/FilamentHelper.py +0 -0
  31. {autoforge-1.5.2/src/autoforge/Helper → autoforge-1.5.4/src/autoforge/Helper/Heightmaps}/__init__.py +0 -0
  32. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Helper/ImageHelper.py +0 -0
  33. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Helper/OutputHelper.py +0 -0
  34. {autoforge-1.5.2/src/autoforge/Loss → autoforge-1.5.4/src/autoforge/Helper}/__init__.py +0 -0
  35. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Loss/PerceptionLoss.py +0 -0
  36. {autoforge-1.5.2/src/autoforge/Misc → autoforge-1.5.4/src/autoforge/Loss}/__init__.py +0 -0
  37. {autoforge-1.5.2/src/autoforge/Modules → autoforge-1.5.4/src/autoforge/Misc}/__init__.py +0 -0
  38. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Misc/color_fit_all.py +0 -0
  39. {autoforge-1.5.2/src/autoforge → autoforge-1.5.4/src/autoforge/Modules}/__init__.py +0 -0
  40. {autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/__main__.py +0 -0
{autoforge-1.5.2 → autoforge-1.5.4}/PKG-INFO RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.2
2
2
  Name: AutoForge
3
- Version: 1.5.2
3
+ Version: 1.5.4
4
4
  Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
5
5
  Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
6
6
  License: CC BY-NC-SA 4.0
@@ -18,6 +18,7 @@ Requires-Dist: pandas>=1.1.0
18
18
  Requires-Dist: scikit-image
19
19
  Requires-Dist: scikit-learn
20
20
  Requires-Dist: tensorboard
21
+ Requires-Dist: pillow>=11.1.0
21
22
 
22
23
  # AutoForge
23
24
 
@@ -26,7 +27,8 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
26
27
  **TLDR:** It uses a picture to generate a 3D layer image that you can print with a 3d printer. Similar to [Hueforge](https://shop.thehueforge.com/), but without the manual work (and without the artistic control).
27
28
 
28
29
  ## Example
29
- All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
30
+ All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge, the background color is set to black.
31
+ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at most 8 different colors and swaps the filament at most 20 times.
30
32
  <div style="display: flex; justify-content: center; gap: 20px;">
31
33
  <div style="text-align: center;">
32
34
  <h3>Input Image</h3>
@@ -37,7 +39,7 @@ All examples use only the 13 BambuLab Basic filaments, currently available in Hu
37
39
  <div style="text-align: center;">
38
40
  <h3>Autoforge Output (75 color layers (3.0mm) + 10 gray background layers (0.4mm)) </h3>
39
41
  <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.png" width="200" />
40
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="200" />
42
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.png" width="200" />
41
43
  <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.png" width="200" />
42
44
  </div>
43
45
  <div style="text-align: center;">
@@ -95,7 +97,7 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
95
97
  - `--iterations`: Number of optimization iterations (default: 5000).
96
98
  - `--learning_rate`: Learning rate for the optimizer (default: 1e-2).
97
99
  - `--layer_height`: Layer thickness in millimeters (default: 0.04).
98
- - `--max_layers`: Maximum number of layers (default: 75). \
100
+ - `--max_layers`: Maximum number of layers (default: 100). \
99
101
  **Note:** This is about 3mm + the background height
100
102
  - `--min_layers`: Minimum number of layers (default: 0). Used to limit height of pruning
101
103
  - `--background_height`: Height of the background in millimeters (default: 0.4). \
@@ -103,18 +105,19 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
103
105
  - `--background_color`: Background color in hexadecimal format (default: `#000000` aka Black). \
104
106
  **Note:** The solver currently assumes that you have a solid color in the background, which means a color with a TD value of 4 or less (if you have a background height of 0.4)
105
107
  - `--output_size`: Maximum dimension for target image (default: 1024).
106
- - `--solver_size`: Maximum dimension for solver (fast) image (default: 128). \
108
+ - `--solver_size`: Maximum dimension for solver (fast) image (default: 256). \
107
109
  **Note:** We solve on a smaller size as this is many times faster, but also a bit less accurate. Increase if you need more accuracy.
108
110
  - `--init_tau`: Initial tau value for Gumbel-Softmax (default: 1.0).
109
111
  - `--final_tau`: Final tau value for the Gumbel-Softmax formulation (default: 0.01).
110
112
  - `--stl_output_size`: Size of the output stl file in mm (default: 200) (20cm).
111
113
  - `--visualize`: Flag to enable live visualization of the composite image during optimization.
112
114
  - `--tensorboard`: Enable TensorBoard logging
113
- - `--run-name`: Name of the run used for TensorBoard logging (optional).
115
+ - `--run_name`: Name of the run used for TensorBoard logging (optional).
114
116
  - `--perform_pruning`: Perform pruning after optimization (default: True). \
115
117
  **Note:** This is highly recommended even if you don't have a color/color swap limit, as it actually increases the quality of the output.
116
118
  - `--pruning_max_colors`: Max number of colors allowed after pruning (default: 100).
117
119
  - `--pruning_max_swaps`: Max number of swaps allowed after pruning (default: 100).
120
+ - `--pruning_max_layers`: Max number of layers allowed after pruning (default: 75).
118
121
  - `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
119
122
 
120
123
  ### Experimental Deph Anything V2 parameters
{autoforge-1.5.2 → autoforge-1.5.4}/README.md RENAMED
@@ -5,7 +5,8 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
5
5
  **TLDR:** It uses a picture to generate a 3D layer image that you can print with a 3d printer. Similar to [Hueforge](https://shop.thehueforge.com/), but without the manual work (and without the artistic control).
6
6
 
7
7
  ## Example
8
- All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
8
+ All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge, the background color is set to black.
9
+ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at most 8 different colors and swaps the filament at most 20 times.
9
10
  <div style="display: flex; justify-content: center; gap: 20px;">
10
11
  <div style="text-align: center;">
11
12
  <h3>Input Image</h3>
@@ -16,7 +17,7 @@ All examples use only the 13 BambuLab Basic filaments, currently available in Hu
16
17
  <div style="text-align: center;">
17
18
  <h3>Autoforge Output (75 color layers (3.0mm) + 10 gray background layers (0.4mm)) </h3>
18
19
  <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.png" width="200" />
19
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="200" />
20
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.png" width="200" />
20
21
  <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.png" width="200" />
21
22
  </div>
22
23
  <div style="text-align: center;">
@@ -74,7 +75,7 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
74
75
  - `--iterations`: Number of optimization iterations (default: 5000).
75
76
  - `--learning_rate`: Learning rate for the optimizer (default: 1e-2).
76
77
  - `--layer_height`: Layer thickness in millimeters (default: 0.04).
77
- - `--max_layers`: Maximum number of layers (default: 75). \
78
+ - `--max_layers`: Maximum number of layers (default: 100). \
78
79
  **Note:** This is about 3mm + the background height
79
80
  - `--min_layers`: Minimum number of layers (default: 0). Used to limit height of pruning
80
81
  - `--background_height`: Height of the background in millimeters (default: 0.4). \
@@ -82,18 +83,19 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
82
83
  - `--background_color`: Background color in hexadecimal format (default: `#000000` aka Black). \
83
84
  **Note:** The solver currently assumes that you have a solid color in the background, which means a color with a TD value of 4 or less (if you have a background height of 0.4)
84
85
  - `--output_size`: Maximum dimension for target image (default: 1024).
85
- - `--solver_size`: Maximum dimension for solver (fast) image (default: 128). \
86
+ - `--solver_size`: Maximum dimension for solver (fast) image (default: 256). \
86
87
  **Note:** We solve on a smaller size as this is many times faster, but also a bit less accurate. Increase if you need more accuracy.
87
88
  - `--init_tau`: Initial tau value for Gumbel-Softmax (default: 1.0).
88
89
  - `--final_tau`: Final tau value for the Gumbel-Softmax formulation (default: 0.01).
89
90
  - `--stl_output_size`: Size of the output stl file in mm (default: 200) (20cm).
90
91
  - `--visualize`: Flag to enable live visualization of the composite image during optimization.
91
92
  - `--tensorboard`: Enable TensorBoard logging
92
- - `--run-name`: Name of the run used for TensorBoard logging (optional).
93
+ - `--run_name`: Name of the run used for TensorBoard logging (optional).
93
94
  - `--perform_pruning`: Perform pruning after optimization (default: True). \
94
95
  **Note:** This is highly recommended even if you don't have a color/color swap limit, as it actually increases the quality of the output.
95
96
  - `--pruning_max_colors`: Max number of colors allowed after pruning (default: 100).
96
97
  - `--pruning_max_swaps`: Max number of swaps allowed after pruning (default: 100).
98
+ - `--pruning_max_layers`: Max number of layers allowed after pruning (default: 75).
97
99
  - `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
98
100
 
99
101
  ### Experimental Deph Anything V2 parameters
{autoforge-1.5.2 → autoforge-1.5.4}/pyproject.toml RENAMED
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "AutoForge"
7
- version="1.5.2"
7
+ version="1.5.4"
8
8
  description = "AutoForge is a Python tool for generating 3D printed layered models from an input image."
9
9
  authors = [{ name = "Hendric Voss", email = "hvoss@techfak.uni-bielefeld.de" }]
10
10
  readme = "README.md"
@@ -22,7 +22,8 @@ dependencies = [
22
22
  "pandas>=1.1.0",
23
23
  "scikit-image",
24
24
  "scikit-learn",
25
- "tensorboard"
25
+ "tensorboard",
26
+ "pillow>=11.1.0"
26
27
  ]
27
28
 
28
29
  [tool.setuptools]
{autoforge-1.5.2 → autoforge-1.5.4}/src/AutoForge.egg-info/PKG-INFO RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.2
2
2
  Name: AutoForge
3
- Version: 1.5.2
3
+ Version: 1.5.4
4
4
  Summary: AutoForge is a Python tool for generating 3D printed layered models from an input image.
5
5
  Author-email: Hendric Voss <hvoss@techfak.uni-bielefeld.de>
6
6
  License: CC BY-NC-SA 4.0
@@ -18,6 +18,7 @@ Requires-Dist: pandas>=1.1.0
18
18
  Requires-Dist: scikit-image
19
19
  Requires-Dist: scikit-learn
20
20
  Requires-Dist: tensorboard
21
+ Requires-Dist: pillow>=11.1.0
21
22
 
22
23
  # AutoForge
23
24
 
@@ -26,7 +27,8 @@ AutoForge is a Python tool for generating 3D printed layered models from an inpu
26
27
  **TLDR:** It uses a picture to generate a 3D layer image that you can print with a 3d printer. Similar to [Hueforge](https://shop.thehueforge.com/), but without the manual work (and without the artistic control).
27
28
 
28
29
  ## Example
29
- All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge.
30
+ All examples use only the 13 BambuLab Basic filaments, currently available in Hueforge, the background color is set to black.
31
+ The pruning is set to a maximum of 8 color and 20 swaps, so each image uses at most 8 different colors and swaps the filament at most 20 times.
30
32
  <div style="display: flex; justify-content: center; gap: 20px;">
31
33
  <div style="text-align: center;">
32
34
  <h3>Input Image</h3>
@@ -37,7 +39,7 @@ All examples use only the 13 BambuLab Basic filaments, currently available in Hu
37
39
  <div style="text-align: center;">
38
40
  <h3>Autoforge Output (75 color layers (3.0mm) + 10 gray background layers (0.4mm)) </h3>
39
41
  <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/lofi_discretized.png" width="200" />
40
- <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.jpg" width="200" />
42
+ <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/nature_discretized.png" width="200" />
41
43
  <img src="https://github.com/hvoss-techfak/AutoForge/blob/main/images/cat_discretized.png" width="200" />
42
44
  </div>
43
45
  <div style="text-align: center;">
@@ -95,7 +97,7 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
95
97
  - `--iterations`: Number of optimization iterations (default: 5000).
96
98
  - `--learning_rate`: Learning rate for the optimizer (default: 1e-2).
97
99
  - `--layer_height`: Layer thickness in millimeters (default: 0.04).
98
- - `--max_layers`: Maximum number of layers (default: 75). \
100
+ - `--max_layers`: Maximum number of layers (default: 100). \
99
101
  **Note:** This is about 3mm + the background height
100
102
  - `--min_layers`: Minimum number of layers (default: 0). Used to limit height of pruning
101
103
  - `--background_height`: Height of the background in millimeters (default: 0.4). \
@@ -103,18 +105,19 @@ autoforge --input_image path/to/input_image.jpg --csv_file path/to/materials.csv
103
105
  - `--background_color`: Background color in hexadecimal format (default: `#000000` aka Black). \
104
106
  **Note:** The solver currently assumes that you have a solid color in the background, which means a color with a TD value of 4 or less (if you have a background height of 0.4)
105
107
  - `--output_size`: Maximum dimension for target image (default: 1024).
106
- - `--solver_size`: Maximum dimension for solver (fast) image (default: 128). \
108
+ - `--solver_size`: Maximum dimension for solver (fast) image (default: 256). \
107
109
  **Note:** We solve on a smaller size as this is many times faster, but also a bit less accurate. Increase if you need more accuracy.
108
110
  - `--init_tau`: Initial tau value for Gumbel-Softmax (default: 1.0).
109
111
  - `--final_tau`: Final tau value for the Gumbel-Softmax formulation (default: 0.01).
110
112
  - `--stl_output_size`: Size of the output stl file in mm (default: 200) (20cm).
111
113
  - `--visualize`: Flag to enable live visualization of the composite image during optimization.
112
114
  - `--tensorboard`: Enable TensorBoard logging
113
- - `--run-name`: Name of the run used for TensorBoard logging (optional).
115
+ - `--run_name`: Name of the run used for TensorBoard logging (optional).
114
116
  - `--perform_pruning`: Perform pruning after optimization (default: True). \
115
117
  **Note:** This is highly recommended even if you don't have a color/color swap limit, as it actually increases the quality of the output.
116
118
  - `--pruning_max_colors`: Max number of colors allowed after pruning (default: 100).
117
119
  - `--pruning_max_swaps`: Max number of swaps allowed after pruning (default: 100).
120
+ - `--pruning_max_layers`: Max number of layers allowed after pruning (default: 75).
118
121
  - `--random_seed`: Random seed for reproducibility (default: 0 (disabled) ).
119
122
 
120
123
  ### Experimental Deph Anything V2 parameters
{autoforge-1.5.2 → autoforge-1.5.4}/src/AutoForge.egg-info/SOURCES.txt RENAMED
@@ -9,7 +9,6 @@ src/AutoForge.egg-info/top_level.txt
9
9
  src/autoforge/__init__.py
10
10
  src/autoforge/__main__.py
11
11
  src/autoforge/auto_forge.py
12
- src/autoforge/old_jax_files.py
13
12
  src/autoforge/Gui/__init__.py
14
13
  src/autoforge/Gui/app.py
15
14
  src/autoforge/Gui/main.py
@@ -24,6 +23,10 @@ src/autoforge/Helper/OutputHelper.py
24
23
  src/autoforge/Helper/PruningHelper.py
25
24
  src/autoforge/Helper/__init__.py
26
25
  src/autoforge/Helper/filamentcolors_library.py
26
+ src/autoforge/Helper/Heightmaps/ChristofidesHeightMap.py
27
+ src/autoforge/Helper/Heightmaps/DepthEstimateHeightMap.py
28
+ src/autoforge/Helper/Heightmaps/OldHeightmap.py
29
+ src/autoforge/Helper/Heightmaps/__init__.py
27
30
  src/autoforge/Loss/LossFunctions.py
28
31
  src/autoforge/Loss/PerceptionLoss.py
29
32
  src/autoforge/Loss/__init__.py
{autoforge-1.5.2 → autoforge-1.5.4}/src/AutoForge.egg-info/requires.txt RENAMED
@@ -10,3 +10,4 @@ pandas>=1.1.0
10
10
  scikit-image
11
11
  scikit-learn
12
12
  tensorboard
13
+ pillow>=11.1.0
{autoforge-1.5.2 → autoforge-1.5.4}/src/autoforge/Helper/CAdamW.py RENAMED
@@ -1,7 +1,6 @@
1
1
  # Taken from https://github.com/kyleliang919/C-Optim/blob/main/c_adamw.py
2
2
  # copy dependencies from transformers/optimization.py
3
3
  import math
4
- import warnings
5
4
  from typing import Callable, Iterable, Tuple
6
5
 
7
6
  import torch
@@ -32,24 +31,33 @@ class CAdamW(Optimizer):
32
31
  """
33
32
 
34
33
  def __init__(
35
- self,
36
- params: Iterable[nn.parameter.Parameter],
37
- lr: float = 1e-3,
38
- betas: Tuple[float, float] = (0.9, 0.999),
39
- eps: float = 1e-6,
40
- weight_decay: float = 0.0,
41
- correct_bias: bool = True,
34
+ self,
35
+ params: Iterable[nn.parameter.Parameter],
36
+ lr: float = 1e-3,
37
+ betas: Tuple[float, float] = (0.9, 0.999),
38
+ eps: float = 1e-6,
39
+ weight_decay: float = 0.0,
40
+ correct_bias: bool = True,
42
41
  ):
43
-
44
42
  if lr < 0.0:
45
43
  raise ValueError(f"Invalid learning rate: {lr} - should be >= 0.0")
46
44
  if not 0.0 <= betas[0] < 1.0:
47
- raise ValueError(f"Invalid beta parameter: {betas[0]} - should be in [0.0, 1.0)")
45
+ raise ValueError(
46
+ f"Invalid beta parameter: {betas[0]} - should be in [0.0, 1.0)"
47
+ )
48
48
  if not 0.0 <= betas[1] < 1.0:
49
- raise ValueError(f"Invalid beta parameter: {betas[1]} - should be in [0.0, 1.0)")
49
+ raise ValueError(
50
+ f"Invalid beta parameter: {betas[1]} - should be in [0.0, 1.0)"
51
+ )
50
52
  if not 0.0 <= eps:
51
53
  raise ValueError(f"Invalid epsilon value: {eps} - should be >= 0.0")
52
- defaults = {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "correct_bias": correct_bias}
54
+ defaults = {
55
+ "lr": lr,
56
+ "betas": betas,
57
+ "eps": eps,
58
+ "weight_decay": weight_decay,
59
+ "correct_bias": correct_bias,
60
+ }
53
61
  super().__init__(params, defaults)
54
62
  self.init_lr = lr
55
63
 
@@ -102,12 +110,16 @@ class CAdamW(Optimizer):
102
110
  if group["correct_bias"]: # No bias correction for Bert
103
111
  bias_correction1 = 1.0 - beta1 ** state["step"]
104
112
  bias_correction2 = 1.0 - beta2 ** state["step"]
105
- step_size = step_size * math.sqrt(bias_correction2) / bias_correction1
113
+ step_size = (
114
+ step_size * math.sqrt(bias_correction2) / bias_correction1
115
+ )
106
116
 
107
117
  # compute norm gradient
108
118
  mask = (exp_avg * grad > 0).to(grad.dtype)
109
119
  # mask = mask * (mask.numel() / (mask.sum() + 1)) ## original implementation, leaving it here for record
110
- mask.div_(mask.mean().clamp_(min=1e-3)) # https://huggingface.co/rwightman/timm-optim-caution found this implementation is more favoarable in many cases
120
+ mask.div_(
121
+ mask.mean().clamp_(min=1e-3)
122
+ ) # https://huggingface.co/rwightman/timm-optim-caution found this implementation is more favoarable in many cases
111
123
  norm_grad = (exp_avg * mask) / denom
112
124
  p.add_(norm_grad, alpha=-step_size)
113
- return loss
125
+ return loss