I defining two model which are encoder and classfication respectively. I want use encoder’s ouput as the input of the classfication then train the two model. After training I want to save the two model separately.
here’s my two model
# encoder model
class source(nn.Module):
def __init__(self, initial_channel):
super(source, self).__init__()
self.initial_channel = initial_channel
def conv_bn(inp, oup, stride):
return nn.Sequential(
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
nn.BatchNorm2d(oup),
nn.ReLU(inplace=True)
)
self.model = self.model = nn.Sequential(conv_bn(initial_channel, 64, 2),
conv_bn(64, 128, 2),
conv_bn(128, 128, 1),
conv_bn(128, 256, 2),
conv_bn(256, 256, 1),
nn.AdaptiveAvgPool2d([1,1]))
def forward(self, x):
x = self.model(x)
x = x.view(-1,256)
return x
# classfication model
class classfication(nn.Module):
def __init__(self, class_num):
super(classfication, self).__init__()
self.class_num = class_num
self.fc1 = nn.Sequential(nn.Linear(256, 64),
nn.ReLU(inplace=True),
nn.Dropout(0.5))
self.fc2 = nn.Sequential(nn.Linear(64, class_num), )
def forward(self, x):
x = self.fc1(x)
x = self.fc2(x)
return x
And following my main model to describe what I want to do in detail.
encoder = source(initial_channel)
cls = classfication(classnum)
model = torch.nn.Sequential(encoder, cls)
train(model)
save(encoder)
save(cls)