polish lr scheduler

python/jittor/__init__.py

@@ -8,7 +8,7 @@
 # This file is subject to the terms and conditions defined in
 # file 'LICENSE.txt', which is part of this source code package.
 # ***************************************************************
-__version__ = '1.2.2.28'
+__version__ = '1.2.2.29'
 from . import lock
 with lock.lock_scope():
     ori_int = int

python/jittor/lr_scheduler.py

@@ -127,7 +127,7 @@ class MultiStepLR(object):
     def get_lr(self):
         now_lr = self.optimizer.lr
         if (self.last_epoch in self.milestones):
-            now_lr *= gamma
+            now_lr *= self.gamma
         return now_lr
 
     def step(self):
@@ -138,4 +138,4 @@ class MultiStepLR(object):
         self.optimizer.lr = new_lr
         for i, param_group in enumerate(self.optimizer.param_groups):
             if param_group.get("lr")!=None:
-                param_group["lr"] = new_lr
+                param_group["lr"] = new_lr

python/jittor/utils/publish.py

@@ -29,6 +29,11 @@ def docker_task(name, build_cmd):
     run_cmd(f"sudo docker save {name}:latest -o /tmp/{bname}.tgz && sudo chmod 666 /tmp/{bname}.tgz")
     upload_file(f"/tmp/{bname}.tgz")
 
+docker_task(
+    "jittor/jittor-cuda-11-1", 
+    "sudo docker build --tag jittor/jittor-cuda-11-1:latest -f script/Dockerfile_cuda11 . --network host"
+)
+
 docker_task(
     "jittor/jittor", 
     "sudo docker build --tag jittor/jittor:latest . --network host"

script/Dockerfile_cuda11

@@ -0,0 +1,51 @@
+# docker build commands
+ARG FROM_IMAGE=nvidia/cuda:11.1.1-cudnn8-devel-ubuntu20.04
+
+FROM ${FROM_IMAGE}
+
+RUN rm /etc/apt/sources.list.d/cuda.list
+
+RUN apt update && apt install ca-certificates -y
+
+# change tsinghua mirror
+RUN echo \
+"deb [trusted=yes] https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ focal main restricted universe multiverse\n\
+deb [trusted=yes] https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ focal-updates main restricted universe multiverse\n\
+deb [trusted=yes] https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ focal-backports main restricted universe multiverse\n\
+deb [trusted=yes] https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ focal-security main restricted universe multiverse" > /etc/apt/sources.list
+
+RUN apt update && apt install wget \
+        python3 python3-dev python3-pip \
+        g++ build-essential -y
+
+WORKDIR /usr/src/jittor
+
+ENV PYTHONIOENCODING utf8
+ENV DEBIAN_FRONTEND noninteractive
+
+# change tsinghua mirror
+RUN pip3 config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
+
+RUN pip3 install  \
+        pybind11 \
+        numpy \
+        tqdm \
+        pillow \
+        astunparse \
+        notebook
+
+RUN pip3 install matplotlib
+
+RUN apt install openmpi-bin openmpi-common libopenmpi-dev -y
+
+RUN pip3 install jittor --timeout 100 && python3 -m jittor.test.test_example
+
+RUN pip3 uninstall jittor -y
+
+COPY . .
+
+RUN pip3 install . --timeout 100
+
+RUN python3 -m jittor.test.test_example
+
+CMD python3 -m jittor.notebook --allow-root --ip=0.0.0.0