The two functions, load_state and load_from_mobilenet, are utility functions for loading pre-trained weights into a PyTorch model. They handle the process of matching the parameters in the checkpoint file with the model's parameters, ensuring compatibility and providing warnings for any mismatches.

1. load_state

def load_state(net, checkpoint):
    source_state = checkpoint['state_dict']
    target_state = net.state_dict()
    new_target_state = collections.OrderedDict()
    for target_key, target_value in target_state.items():
        if target_key in source_state and source_state[target_key].size() == target_state[target_key].size():
            new_target_state[target_key] = source_state[target_key]
        else:
            new_target_state[target_key] = target_state[target_key]
            print('[WARNING] Not found pre-trained parameters for {}'.format(target_key))

    net.load_state_dict(new_target_state)

Purpose:

• Loads a pre-trained state dictionary (state_dict) from a checkpoint into the model (net).
• Ensures that only parameters with matching names and sizes are loaded.
• Provides warnings for parameters in the model that do not have corresponding pre-trained weights.

How It Works:

1. Extract State Dictionaries:
   • source_state: The state dictionary from the checkpoint.
   • target_state: The state dictionary of the model (net).
2. Iterate Over Model Parameters:
   • For each parameter in the model (target_key):
     • If the parameter exists in the checkpoint (source_state) and its size matches, it is loaded from the checkpoint.
     • Otherwise, the model's default parameter is retained, and a warning is printed.
3. Load the Updated State Dictionary:
   • The updated state dictionary (new_target_state) is loaded into the model using net.load_state_dict.

Use Case:

• General-purpose function for loading pre-trained weights into a model, ensuring compatibility and handling missing or mismatched parameters gracefully.

2. load_from_mobilenet

def load_from_mobilenet(net, checkpoint):
    source_state = checkpoint['state_dict']
    target_state = net.state_dict()
    new_target_state = collections.OrderedDict()
    for target_key, target_value in target_state.items():
        k = target_key
        if k.find('model') != -1:
            k = k.replace('model', 'module.model')
        if k in source_state and source_state[k].size() == target_state[target_key].size():
            new_target_state[target_key] = source_state[k]
        else:
            new_target_state[target_key] = target_state[target_key]
            print('[WARNING] Not found pre-trained parameters for {}'.format(target_key))

    net.load_state_dict(new_target_state)

Purpose:

• Specifically designed to load pre-trained weights from a MobileNet-based checkpoint into a model.
• Handles cases where the parameter names in the checkpoint differ slightly from those in the model (e.g., due to the use of DataParallel).