I see, one more question, will you move the model around, such as on a different machine with different gpu number, or are you loading the whole model on the same devices?
If you don’t, and you really want to save them seperately to different files, maybe for better inspection or archive perpose, then:
If you do, then you will have to decide which device each part of your model would locate on, eg: suppose on your training machine you have 3 gpus, and on your inference machine you have 1 gpu.
your_model.fc1 = torch.load("fc1.pt", map_location=torch.device('cuda:0'))
your_model.fc2 = torch.load("fc2.pt", map_location=torch.device('cuda:0'))
your_model.fc3 = torch.load("fc3.pt", map_location=torch.device('cuda:0'))
by the way,
Maybe you have some wrong idea, there is not such a “connected device” concept in pytorch, you can perform a complex forward() operation or a simple add() operation on some input
x locating on device
cpu simply because the operands (tensors) locates on the same device, if torch needs to fetch it somewhere else, it will complain and throw an error.
About saving the model
There are many ways to save your model, typically you will want to save the
OrderedDict returned by
model.state_dict(), the keys are your parameter names such as “linear.weight” or “linear.bias”, and values are
.data attribute is just a Tensor. You may load a state dict into your model like:
def prep_load_state_dict(model: nn.Module,
Automatically load a **loaded state dictionary**
This function handles tensor device remapping.
for name, param in model.named_parameters():
About torch.save and torch.load
If you know the
pickle concept in python, then you will get what
pickle serialize a object into binary string:
buffer = io.BytesIO()
you can serialize whatever you like into this, cuda tensor will essentially be saved as “raw data” + “device descriptor cuda:0”.