Pytorch自定义dataloader以及在迭代过程中返回image的(3)


        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

def _make_layer(self, block, planes, blocks, stride=1):
        #  block: object, planes: output channel, blocks: the num of blocks
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion  # the input channel num become 4 times
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

return nn.Sequential(*layers)

def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x


def resnet50(pretrained = True):
    """Constructs a ResNet-50 model.

Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model = ResNet(Bottleneck, [3, 4, 6, 3])
    # model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
    model.load_state_dict(torch.load('./resnet50_20170907_state_dict.pth'))
    return model
cnn = resnet50(pretrained=True)  # the output number is 9
cnn.cuda()
cnn.eval()
criterion = nn.MSELoss().cuda()

for i, (test_images, test_labels, fn) in enumerate(test_loader):  # the first i in index, and the () is the content
    test_images = Variable(test_images.cuda())
    test_labels = Variable(test_labels.cuda())
    outputs = cnn(test_images)
    print(outputs.data[0])
    print(fn)
    loss = criterion(outputs, test_labels)
    print("Iter [%d/%d] Test_Loss: %.4f" % (i + 1, 781, loss.data[0]))

着重看定义dataloader以及返回图像名称的一段代码:

def random_choose_data(label_path):
    random.seed(1)
    file = open(label_path)
    lines = file.readlines()
    slice_initial = random.sample(lines, 200000)  # if don't change this ,it will be all the same
    slice = list(set(lines)-set(slice_initial))
    random.shuffle(slice)

train_label = slice[:150000]
    test_label = slice[150000:200000]
    return train_label, test_label  # output the list and delvery it into ImageFolder

内容版权声明:除非注明,否则皆为本站原创文章。

转载注明出处:https://www.heiqu.com/1f2de4ccea99914fdd992d479ea496a6.html