# 載入雲端硬碟 
from google.colab import drive 
drive.mount('/content/drive') 

# 查看分配到的GPU 
gpu_info = !nvidia-smi 
gpu_info = '\n'.join(gpu_info) 
if gpu_info.find('failed') >= 0: 
    print('Not connected to a GPU') 
else: 
    print(gpu_info) 

# 安裝python包 
!pip3 install <package>

# 克隆倉庫到/content/my-repo目錄下 
!git clone //github.com/my-github-username/my-git-repo.git %cd my-git-repo !./train.py --logdir /my/log/path --data_root /my/data/root --resume

from train import my_training_method 
my_training_method(arg1, arg2, ...)

import sys 
sys.path.append('/content/my-git-repo') # 把git倉庫的目錄添加到系統目錄

from google.colab import drive 
drive.mount('/content/drive')

!python3 /content/drive/MyDrive/zhihu_colab/mldl_project/code/train_eval.py \ 
--data_root /content/datasets/ROD-synROD \ 
--logdir /content/drive/MyDrive/ \ --
resume \ 
| tee /content/drive/MyDrive/synRODtoROD.txt -a

def save_checkpoint(path: Text,
                    epoch: int,
                    modules: Union[nn.Module, Sequence[nn.Module]],
                    optimizers: Union[opt.Optimizer, Sequence[opt.Optimizer]],
                    safe_replacement: bool = True):
    """
    Save a checkpoint of the current state of the training, so it can be resumed.
    This checkpointing function assumes that there are no learning rate schedulers or gradient scalers for automatic
    mixed precision.
    :param path:
        Path for your checkpoint file
    :param epoch:
        Current (completed) epoch
    :param modules:
        nn.Module containing the model or a list of nn.Module objects
    :param optimizers:
        Optimizer or list of optimizers
    :param safe_replacement:
        Keep old checkpoint until the new one has been completed
    :return:
    """

    # This function can be called both as
    # save_checkpoint('/my/checkpoint/path.pth', my_epoch, my_module, my_opt)
    # or
    # save_checkpoint('/my/checkpoint/path.pth', my_epoch, [my_module1, my_module2], [my_opt1, my_opt2])
    if isinstance(modules, nn.Module):
        modules = [modules]
    if isinstance(optimizers, opt.Optimizer):
        optimizers = [optimizers]
 
    # Data dictionary to be saved
    data = {
        'epoch': epoch,
        # Current time (UNIX timestamp)
        'time': time.time(),
        # State dict for all the modules
        'modules': [m.state_dict() for m in modules],
        # State dict for all the optimizers
        'optimizers': [o.state_dict() for o in optimizers]
    }

    # Safe replacement of old checkpoint
    temp_file = None
    if os.path.exists(path) and safe_replacement:
        # There's an old checkpoint. Rename it!
        temp_file = path + '.old'
        os.rename(path, temp_file)

    # Save the new checkpoint
    with open(path, 'wb') as fp:
        torch.save(data, fp)
        # Flush and sync the FS
        fp.flush()
        os.fsync(fp.fileno())

    # Remove the old checkpoint
    if temp_file is not None:
        os.unlink(path + '.old')


def load_checkpoint(path: Text,
                    default_epoch: int,
                    modules: Union[nn.Module, Sequence[nn.Module]],
                    optimizers: Union[opt.Optimizer, Sequence[opt.Optimizer]],
                    verbose: bool = True):
    """
    Try to load a checkpoint to resume the training.
    :param path:
        Path for your checkpoint file
    :param default_epoch:
        Initial value for "epoch" (in case there are not snapshots)
    :param modules:
        nn.Module containing the model or a list of nn.Module objects. They are assumed to stay on the same device
    :param optimizers:
        Optimizer or list of optimizers
    :param verbose:
        Verbose mode
    :return:
        Next epoch
    """
    if isinstance(modules, nn.Module):
        modules = [modules]
    if isinstance(optimizers, opt.Optimizer):
        optimizers = [optimizers]

    # If there's a checkpoint
    if os.path.exists(path):
        # Load data
        data = torch.load(path, map_location=next(modules[0].parameters()).device)

        # Inform the user that we are loading the checkpoint
        if verbose:
            print(f"Loaded checkpoint saved at {datetime.fromtimestamp(data['time']).strftime('%Y-%m-%d %H:%M:%S')}. "
                  f"Resuming from epoch {data['epoch']}")

        # Load state for all the modules
        for i, m in enumerate(modules):
            modules[i].load_state_dict(data['modules'][i])

        # Load state for all the optimizers
        for i, o in enumerate(optimizers):
            optimizers[i].load_state_dict(data['optimizers'][i])

        # Next epoch
        return data['epoch'] + 1
    else:
        return default_epoch

if args.resume: # args.resume是命令行輸入的參數，用於指示要不要載入上次訓練的結果 
    first_epoch = load_checkpoint(checkpoint_path, first_epoch, net_list, optims_list)

# Save checkpoint 
save_checkpoint(checkpoint_path, epoch, net_list, optims_list)