# @Time : 2022/1/19 17:57 # @Author : PeinuanQin# @File : test.pyimport torch.nn as nn
import torch
import torch.utils.data as Data
from tqdm import tqdm
from torchvision import transforms,datasets
import numpy as np
import torchvision
from torch.optim import lr_scheduler
classA(nn.Module):def__init__(self):super(A,self).__init__()
self.conv = nn.Conv2d(in_channels=3,out_channels=8,kernel_size=3)
self.relu = nn.ReLU(inplace=True)defforward(self,x):
out = self.conv(x)
out = self.relu(out)
B_model = B()
out = B_model(out)return out
classB(nn.Module):def__init__(self):super(B,self).__init__()
self.conv = nn.Conv2d(in_channels=8,out_channels=16,kernel_size=3)
self.relu = nn.ReLU(inplace=True)defforward(self, x):
out = self.conv(x)
out = self.relu(out)return out
这个时候就会报错,而报错的原因,就是因为 torch 的流程是这样的:
首先将所有的模型加载,先从 A 开始,进入 A 的 init 中把所有的内容加载,然后,通过 main 函数中的 to(device) 操作,就把加载的所有内容和网络定义都放到 cuda 上了,但是注意!!!
第二步开始训练,训练的过程中,都是通过 forward 函数来调用的,但是这个时候程序发现,当进入 A 的 forward 中运行的时候,出现了几个 B 的网络层,但是注意:这些 B 中定义的网络层,在网络加载的过程中可是没有出现在 A 的 __init__里面,也就理所当然地没有加载到 cuda上,因此在 A 的 forward 中出现的时候,B 的这几个网络层的 weight 依然在 cpu 上,这就导致了错误。
classA(nn.Module):def__init__(self):super(A,self).__init__()
self.conv = nn.Conv2d(in_channels=3,out_channels=8,kernel_size=3)
self.relu = nn.ReLU(inplace=True)
self.b_module = B()defforward(self,x):
out = self.conv(x)
out = self.relu(out)
out = self.b_module(out)return out
classB(nn.Module):def__init__(self):super(B,self).__init__()
self.conv = nn.Conv2d(in_channels=8,out_channels=16,kernel_size=3)
self.relu = nn.ReLU(inplace=True)defforward(self, x):
out = self.conv(x)
out = self.relu(out)return out