[NFC] polish colossalai/gemini/gemini_context.py code style (#2690)

colossalai/gemini/gemini_context.py

@@ -1,48 +1,48 @@
-from enum import EnumMeta
+from enum import EnumMeta
-
+
-
+
-class GeminiMemoryManager(object):
+class GeminiMemoryManager(object):
-
+
-    def __init__(self, states_cls: EnumMeta):
+    def __init__(self, states_cls: EnumMeta):
-        super().__init__()
+        super().__init__()
-        self.states_cls = states_cls
+        self.states_cls = states_cls
-        self._cnter = 0    # the counter of instances
+        self._cnter = 0    # the counter of instances
-
+
-        self.total_mem = dict()
+        self.total_mem = dict()
-        self.state_mem = dict()
+        self.state_mem = dict()
-        self.state_mem['cpu'] = dict()
+        self.state_mem['cpu'] = dict()
-        self.state_mem['cuda'] = dict()
+        self.state_mem['cuda'] = dict()
-
+
-        self.reset()
+        self.reset()
-
+
-    @property
+    @property
-    def total_number(self):
+    def total_number(self):
-        return self._cnter
+        return self._cnter
-
+
-    def reset(self):
+    def reset(self):
-        self._cnter = 0    # the counter of instances
+        self._cnter = 0    # the counter of instances
-
+
-        self.total_mem['cpu'] = 0    # memory occupation of instances in cpu
+        self.total_mem['cpu'] = 0    # memory occupation of instances in cpu
-        self.total_mem['cuda'] = 0    # memory of occupation of instances in cuda
+        self.total_mem['cuda'] = 0    # memory of occupation of instances in cuda
-
+
-        # memory conditions for all states
+        # memory conditions for all states
-        for state in self.states_cls:
+        for state in self.states_cls:
-            self.state_mem['cpu'][state] = 0
+            self.state_mem['cpu'][state] = 0
-            self.state_mem['cuda'][state] = 0
+            self.state_mem['cuda'][state] = 0
-
+
-    def register_new_instance(self):
+    def register_new_instance(self):
-        self._cnter += 1
+        self._cnter += 1
-
+
-    def delete_instance(self):
+    def delete_instance(self):
-        self._cnter -= 1
+        self._cnter -= 1
-
+
-    def print_info(self):
+    def print_info(self):
-        print(f"Total number: {self.total_number}",
+        print(f"Total number: {self.total_number}",
-              f"Total CPU memory occupation: {self.total_mem['cpu']}",
+              f"Total CPU memory occupation: {self.total_mem['cpu']}",
-              f"Total CUDA memory occupation: {self.total_mem['cuda']}\n",
+              f"Total CUDA memory occupation: {self.total_mem['cuda']}\n",
-              sep='\n')
+              sep='\n')
-
+
-        for state in self.states_cls:
+        for state in self.states_cls:
-            print(f"{state}: CPU memory occupation: {self.state_mem['cpu'][state]}",
+            print(f"{state}: CPU memory occupation: {self.state_mem['cpu'][state]}",
-                  f"{state}: CUDA memory occupation: {self.state_mem['cuda'][state]}\n",
+                  f"{state}: CUDA memory occupation: {self.state_mem['cuda'][state]}\n",
-                  sep='\n')
+                  sep='\n')